FLUX Fine-Tuning for Gaudi

[dsocek]

FLUX Fine-Tuning for Gaudi

Overview

FLUX.1-dev is a high-quality text-to-image model that has been attracting significant attention for its impressive image generation capabilities. In this PR, we introduce support for FLUX fine-tuning on Gaudi using the LoRA DreamBooth adapter. While the model is highly memory-intensive for training, we've developed a training script specifically for Gaudi devices that fits within HPU memory constraints and enables effective fine-tuning. The implementation works on a single Gaudi device as well as a multi-card Gaudi system (supporting both MPI and DeepSpeed training).

List of Contributions in this PR

• FLUX.1-dev LoRA DreamBooth training script for Gaudi
• Single-card and multi-card (MPI and DeepSpeed) support
• Updated documentation with FLUX training samples
• Training FLUX CI test added diffusers tests

Quality

Ran several quality tests, and FLUX fine-tuning script demonstrated good visual results. The following example was ran on Gaudi2.

Training set ("dog" dataset of 5 images from Hugging Face):

Ran fine-tuning of FLUX.1-dev model on Gaudi2 with 1000 training steps and captured checkpoints after every 250 steps. The following 5 images are ran with prompt "A photo of sks dog in a bucket" using original FLUX.1-dev model and then 4 fine-tuned checkpoints. The images show a clear progression in the dog's resemblance to the one from the training set:

Performance

Both single card and multi-card FLUX training on Gaudi2 was evaluated for performance.

Single card performance

Ran single card FLUX.1-dev training on Gaudi2 with:

python train_dreambooth_lora_flux.py \
  --pretrained_model_name_or_path="black-forest-labs/FLUX.1-dev" \
  --dataset="dog" \
  --prompt="a photo of sks dog" \
  --output_dir="dog_lora_flux" \
  --mixed_precision="bf16" \
  --weighting_scheme="none" \
  --resolution=1024 \
  --train_batch_size=1 \
  --learning_rate=1e-4 \
  --guidance_scale=1 \
  --report_to="tensorboard" \
  --gradient_accumulation_steps=4 \
  --gradient_checkpointing \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --cache_latents \
  --rank=4 \
  --max_train_steps=500 \
  --seed="0" \
  --use_hpu_graphs_for_inference \
  --use_hpu_graphs_for_training \
  --gaudi_config_name="Habana/stable-diffusion"

Output:

11/11/2024 16:07:42 - INFO - __main__ - ***** Running training *****
11/11/2024 16:07:42 - INFO - __main__ -   Num examples = 5
11/11/2024 16:07:42 - INFO - __main__ -   Num batches each epoch = 5
11/11/2024 16:07:42 - INFO - __main__ -   Num Epochs = 250
11/11/2024 16:07:42 - INFO - __main__ -   Instantaneous batch size per device = 1
11/11/2024 16:07:42 - INFO - __main__ -   Total train batch size (w. parallel, distributed & accumulation) = 4
11/11/2024 16:07:42 - INFO - __main__ -   Gradient Accumulation steps = 4
11/11/2024 16:07:42 - INFO - __main__ -   Total optimization steps = 500
Steps: 100%|█████████████████████| 500/500 [1:53:10<00:00, 11.50s/it, loss=0.165, lr=0.0001]
Model weights saved in dog_lora_flux/pytorch_lora_weights.safetensors

Multi-card performance

Ran 8 card FLUX.1-dev training on Gaudi2 using MPI with:

python ../../gaudi_spawn.py --world_size 8 --use_mpi train_dreambooth_lora_flux.py \
  --pretrained_model_name_or_path="black-forest-labs/FLUX.1-dev" \
  --dataset="dog" \
  --prompt="a photo of sks dog" \
  --output_dir="dog_lora_flux" \
  --mixed_precision="bf16" \
  --weighting_scheme="none" \
  --resolution=1024 \
  --train_batch_size=1 \
  --learning_rate=1e-4 \
  --guidance_scale=1 \
  --report_to="tensorboard" \
  --gradient_accumulation_steps=4 \
  --gradient_checkpointing \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --cache_latents \
  --rank=4 \
  --max_train_steps=500 \
  --seed="0" \
  --use_hpu_graphs_for_inference \
  --use_hpu_graphs_for_training \
  --gaudi_config_name="Habana/stable-diffusion"

Output:

11/11/2024 17:10:52 - INFO - __main__ - ***** Running training *****
11/11/2024 17:10:52 - INFO - __main__ -   Num examples = 5
11/11/2024 17:10:52 - INFO - __main__ -   Num batches each epoch = 1
11/11/2024 17:10:52 - INFO - __main__ -   Num Epochs = 500
11/11/2024 17:10:52 - INFO - __main__ -   Instantaneous batch size per device = 1
11/11/2024 17:10:52 - INFO - __main__ -   Total train batch size (w. parallel, distributed & accumulation) = 32
11/11/2024 17:10:52 - INFO - __main__ -   Gradient Accumulation steps = 4
11/11/2024 17:10:52 - INFO - __main__ -   Total optimization steps = 500
Steps: 100%|██████████| 500/500 [46:09<00:00,  5.13s/it, loss=0.74, lr=0.0001]
Model weights saved in dog_lora_flux/pytorch_lora_weights.safetensors

Ran 8 card FLUX.1-dev training on Gaudi2 using DeepSpeed with:

python ../../gaudi_spawn.py --world_size 8 --use_deepspeed train_dreambooth_lora_flux.py \
  --pretrained_model_name_or_path="black-forest-labs/FLUX.1-dev" \
  --dataset="dog" \
  --prompt="a photo of sks dog" \
  --output_dir="dog_lora_flux" \
  --mixed_precision="bf16" \
  --weighting_scheme="none" \
  --resolution=1024 \
  --train_batch_size=1 \
  --learning_rate=1e-4 \
  --guidance_scale=1 \
  --report_to="tensorboard" \
  --gradient_accumulation_steps=4 \
  --gradient_checkpointing \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --cache_latents \
  --rank=4 \
  --max_train_steps=500 \
  --seed="0" \
  --use_hpu_graphs_for_inference \
  --use_hpu_graphs_for_training \
  --gaudi_config_name="Habana/stable-diffusion"

Output:

11/11/2024 15:27:54 - INFO - __main__ - ***** Running training *****
11/11/2024 15:27:54 - INFO - __main__ -   Num examples = 5
11/11/2024 15:27:54 - INFO - __main__ -   Num batches each epoch = 1
11/11/2024 15:27:54 - INFO - __main__ -   Num Epochs = 500
11/11/2024 15:27:54 - INFO - __main__ -   Instantaneous batch size per device = 1
11/11/2024 15:27:54 - INFO - __main__ -   Total train batch size (w. parallel, distributed & accumulation) = 32
11/11/2024 15:27:54 - INFO - __main__ -   Gradient Accumulation steps = 4
11/11/2024 15:27:54 - INFO - __main__ -   Total optimization steps = 500
Steps: 100%|██████████| 500/500 [45:58<00:00,  5.11s/it, loss=0.74, lr=0.0001]
Model weights saved in dog_lora_flux/pytorch_lora_weights.safetensors

Device	Configuration	Synapse	Training Throughput [s/it]
Gaudi2	1-HPU	1.18.0-524	11.50
Gaudi2	8-HPU w/ MPI	1.18.0-524	5.13
Gaudi2	8-HPU w/ DeepSpeed	1.18.0-524	5.11

Tests

Added CI test for FLUX training PASSED ✔️

python -m pytest tests/test_diffusers.py -v -s -k "test_dreambooth_lora_flux"
...
============ 1 passed, 162 deselected, 1 warning in 18.61s =============

[imangohari1]

@dsocek
Hi Daniel,
Thanks for adding these. I am working on a deeper review and testing on this but added some suggestions.

I also am not sure why we have these empty files? Maybe they need to be deleted?

Empty file modified 0
[examples/stable-diffusion/training/train_dreambooth_lora_sdxl.py](https://github.com/huggingface/optimum-habana/pull/1482/files?diff=unified&w=0#diff-f70e088d9b743c02b0e00a1de796c1397b19cd509d586962a403bddafc19af3d)
100644 → 100755
Viewed
Empty file.
Empty file modified 0
[examples/stable-diffusion/training/train_text_to_image_sdxl.py](https://github.com/huggingface/optimum-habana/pull/1482/files?diff=unified&w=0#diff-7a22acb0f64524f36bb39e1c9b6b9c192abd3cef9ebc1d4b4a5e0169dcd8a416)
100644 → 100755
Viewed
Empty file.
Empty file modified 0
[examples/stable-diffusion/unconditional_image_generation.py](https://github.com/huggingface/optimum-habana/pull/1482/files?diff=unified&w=0#diff-4e3a2db778ea0b3c6724c37248c31326e50fd7eb8351a2130a91ca288c78b9fa)
100644 → 100755
Viewed
Empty file.

[imangohari1]

should we bring this to its own simple few line script download_data_dog.py and then here have it run like python XXX?
The mix of bash and python here is confusing.

[dsocek]

I wanted to make it cut and paste ready. TBH this is a matter of preference... I am not in favor of adding download scripts as files. We also have cat dataset as example for prior training samples. Actually now there is no python, all sample code is bash :)

[regisss]

I agree that it's easier to use that way, but on the other hand it makes the code snippet much less clear IMO, see for example this one: https://github.com/dsocek/optimum-habana/tree/flux-fine-tuning/examples/stable-diffusion/training#cat-toy-example
I would rather keep pure Python code snippets.

[imangohari1]

same here.
should we bring this to its own simple few line script download_data_cat.py and then here have it run like python XXX?
The mix of bash and python here is confusing.

[regisss]

I agree with @imangohari1, why using python -c "..." instead of simply having a Python script?
edit: I saw the reply below, continuing there

[dsocek]

@dsocek Hi Daniel, Thanks for adding these. I am working on a deeper review and testing on this but added some suggestions.

I also am not sure why we have these empty files? Maybe they need to be deleted?

Empty file modified 0
[examples/stable-diffusion/training/train_dreambooth_lora_sdxl.py](https://github.com/huggingface/optimum-habana/pull/1482/files?diff=unified&w=0#diff-f70e088d9b743c02b0e00a1de796c1397b19cd509d586962a403bddafc19af3d)
100644 → 100755
Viewed
Empty file.
Empty file modified 0
[examples/stable-diffusion/training/train_text_to_image_sdxl.py](https://github.com/huggingface/optimum-habana/pull/1482/files?diff=unified&w=0#diff-7a22acb0f64524f36bb39e1c9b6b9c192abd3cef9ebc1d4b4a5e0169dcd8a416)
100644 → 100755
Viewed
Empty file.
Empty file modified 0
[examples/stable-diffusion/unconditional_image_generation.py](https://github.com/huggingface/optimum-habana/pull/1482/files?diff=unified&w=0#diff-4e3a2db778ea0b3c6724c37248c31326e50fd7eb8351a2130a91ca288c78b9fa)
100644 → 100755
Viewed
Empty file.

@imangohari1 I changed file permissions from 644 (r+w) to 755 (added x) so people can just execute in bash without having to put python .. This is better. Also half files were 644 and half were 755. Adds consistency and better UX

[imangohari1]

Hi Daniel,
I did more testing on this and here are some suggested changes:

• There are some typos/minor-issues in README. the flux pytest does not run on G1 since it gives an OOM. I fixed them all in the attached patch. Please apply it with git am < 0001*

  • 0001-fea-Added-some-fixed.-Added-G1-pytest-for-flux.patch

• When using the flux fine-tuned data with text_to_image_generation.py I am seeing an error ValueError: The current scheduler class <class 'optimum.habana.diffusers.schedulers.scheduling_ddim.GaudiDDIMScheduler'>'s set_timesteps does not support custom sigmas schedules. Please check whether you are using the correct scheduler. (details below). Can you please investigate this?

testing details

fine-tuning using mpi on 8x.

python ../../gaudi_spawn.py --world_size 8 --use_mpi train_dreambooth_lora_flux.py   --pretrained_model_name_or_path="black-forest-labs/FLUX.1-dev"   --dataset="dog"   --prompt="a photo of sks dog"   --output_dir="dog_lora_flux"   --mixed_precision="bf16"   --weighting_scheme="none"   --resolution=1024   --train_batch_size=1   --learning_rate=1e-4   --guidance_scale=1   --report_to="tensorboard"   --gradient_accumulation_steps=4   --gradient_checkpointing   --lr_scheduler="constant"   --lr_warmup_steps=0   --cache_latents   --rank=4   --max_train_steps=500   --seed="0"   --use_hpu_graphs_for_inference   --use_hpu_graphs_for_training   --gaudi_config_name="Habana/stable-diffusion"

runs fine and then

python ../text_to_image_generation.py     --model_name_or_path "black-forest-labs/FLUX.1-dev"     --lora_id dog_lora_flux     --prompts "A picture of a sks dog in a bucket"     --num_images_per_prompt 5     --batch_size 1     --image_save_dir /tmp/flux_images     --use_habana     --use_hpu_graphs     --gaudi_config Habana/stable-diffusion     --bf16

crashes with

[WARNING|pipeline_utils.py:157] 2024-11-19 20:57:41,645 >> `use_torch_autocast` is True in the given Gaudi configuration but `torch_dtype=torch.bfloat16` was given. Disabling mixed precision and continuing in bf16 only.
============================= HABANA PT BRIDGE CONFIGURATION ===========================
 PT_HPU_LAZY_MODE = 1
 PT_RECIPE_CACHE_PATH =
 PT_CACHE_FOLDER_DELETE = 0
 PT_HPU_RECIPE_CACHE_CONFIG =
 PT_HPU_MAX_COMPOUND_OP_SIZE = 9223372036854775807
 PT_HPU_LAZY_ACC_PAR_MODE = 1
 PT_HPU_ENABLE_REFINE_DYNAMIC_SHAPES = 0
 PT_HPU_EAGER_PIPELINE_ENABLE = 1
 PT_HPU_EAGER_COLLECTIVE_PIPELINE_ENABLE = 1
---------------------------: System Configuration :---------------------------
Num CPU Cores : 160
CPU RAM       : 1056374404 KB
------------------------------------------------------------------------------
Traceback (most recent call last):
  File "/devops/sgohari/tests/codes/pr-reviews/1482/optimum-habana/examples/stable-diffusion/training/../text_to_image_generation.py", line 644, in <module>
    main()
  File "/devops/sgohari/tests/codes/pr-reviews/1482/optimum-habana/examples/stable-diffusion/training/../text_to_image_generation.py", line 613, in main
    outputs = pipeline(prompt=args.prompts, **kwargs_call)
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/optimum/habana/diffusers/pipelines/flux/pipeline_flux.py", line 589, in __call__
    timesteps, num_inference_steps = retrieve_timesteps(
  File "/usr/local/lib/python3.10/dist-packages/diffusers/pipelines/flux/pipeline_flux.py", line 127, in retrieve_timesteps
    raise ValueError(
ValueError: The current scheduler class <class 'optimum.habana.diffusers.schedulers.scheduling_ddim.GaudiDDIMScheduler'>'s `set_timesteps` does not support custom sigmas schedules. Please check whether you are using the correct scheduler.

[dsocek]

@imangohari1 thanks for minor issue find, I fixed in amend should be good now

[imangohari1]

LGTM! @regisss please review.

[dsocek]

rebased

[regisss]

I left a few comments.

Also, I realize that we didn't add FLUX to the table of validated diffusion models in the README and in the docs, can you do it in this PR please?

[regisss]

I agree with @imangohari1, why using python -c "..." instead of simply having a Python script?
edit: I saw the reply below, continuing there

[regisss]

I agree that it's easier to use that way, but on the other hand it makes the code snippet much less clear IMO, see for example this one: https://github.com/dsocek/optimum-habana/tree/flux-fine-tuning/examples/stable-diffusion/training#cat-toy-example
I would rather keep pure Python code snippets.

[dsocek]

@regisss all fixed and added FLUX.1 to main README and docs/source/index

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[the-pikachu]

Single card performance

Ran single card FLUX.1-dev training on Gaudi2 with:

python train_dreambooth_lora_flux.py \
  --pretrained_model_name_or_path="black-forest-labs/FLUX.1-dev" \
  --dataset="dog" \
  --prompt="a photo of sks dog" \
  --output_dir="dog_lora_flux" \
  --mixed_precision="bf16" \
  --weighting_scheme="none" \
  --resolution=1024 \
  --train_batch_size=1 \
  --learning_rate=1e-4 \
  --guidance_scale=1 \
  --report_to="tensorboard" \
  --gradient_accumulation_steps=4 \
  --gradient_checkpointing \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --cache_latents \
  --rank=4 \
  --max_train_steps=500 \
  --seed="0" \
  --use_hpu_graphs_for_inference \
  --use_hpu_graphs_for_training \
  --gaudi_config_name="Habana/stable-diffusion"

Output:

11/11/2024 16:07:42 - INFO - __main__ - ***** Running training *****
11/11/2024 16:07:42 - INFO - __main__ -   Num examples = 5
11/11/2024 16:07:42 - INFO - __main__ -   Num batches each epoch = 5
11/11/2024 16:07:42 - INFO - __main__ -   Num Epochs = 250
11/11/2024 16:07:42 - INFO - __main__ -   Instantaneous batch size per device = 1
11/11/2024 16:07:42 - INFO - __main__ -   Total train batch size (w. parallel, distributed & accumulation) = 4
11/11/2024 16:07:42 - INFO - __main__ -   Gradient Accumulation steps = 4
11/11/2024 16:07:42 - INFO - __main__ -   Total optimization steps = 500
Steps: 100%|█████████████████████| 500/500 [1:53:10<00:00, 11.50s/it, loss=0.165, lr=0.0001]
Model weights saved in dog_lora_flux/pytorch_lora_weights.safetensors

Hi Team,

Actually i runned the training code(train_dreambooth_flux_lora.py) provided for flux model , what exactly in the above, the same only i runned with my images (13 images) and the prompt. It was runned on Gaudi2 machine with single card. but the problem is its not capturing the images features when i inference with the trained safetensors. what is the issue here? I am unable to figure out the issue. Please suggest the way to train with my image features. I did not crop the images, all quality images only.