How to finetune Flux-dev full params, 80G OOM ...

[huangjun12]

I am using the train_dreambooth_flux script to fine-tune the flux-dev model with full parameters using DeepSpeed Stage 2. However, I am still encountering out-of-memory issues on an 80GB GPU. Are there any solutions available to address this problem? Thanks!

[hlky]

cc @linoytsaban @sayakpaul

[sayakpaul]

What is the training command you're using?

[huangjun12]

What is the training command you're using?

train_dreambooth_flux

export MODEL_NAME="black-forest-labs/FLUX.1-dev"
export INSTANCE_DIR="dog"
export OUTPUT_DIR="trained-flux"

accelerate launch train_dreambooth_flux.py \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --instance_data_dir=$INSTANCE_DIR \
  --output_dir=$OUTPUT_DIR \
  --mixed_precision="bf16" \
  --instance_prompt="a photo of sks dog" \
  --resolution=1024 \
  --train_batch_size=1 \
  --guidance_scale=1 \
  --gradient_accumulation_steps=4 \
  --optimizer="prodigy" \
  --learning_rate=1. \
  --report_to="wandb" \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=500 \
  --validation_prompt="A photo of sks dog in a bucket" \
  --validation_epochs=25 \
  --seed="0" \
  --push_to_hub

as in README_flux.md

[sayakpaul]

Can you try:

--gradient_checkpointing
--use_8bit_adam

[hlky]

Could be validation as in #10145

[huangjun12]

Can you try:

--gradient_checkpointing
--use_8bit_adam

I tried this, and turn off validation, it still OOM...

[sayakpaul]

@linoytsaban can you give this a look?

[huangjun12]

Can you try:

--gradient_checkpointing
--use_8bit_adam

Have you successfully fine-tuned the full flux-dev model using 80GB GPU memory? What are the environment dependencies?

[timmerscher]

I'm encountering the same issue!
While tracking it with htop the memory-usage rises above 100GB

Error message is:

Traceback (most recent call last):
  File "/home/tim/miniconda3/envs/flux/bin/accelerate", line 11, in <module>
    sys.exit(main())
             ^^^^^^
  File "/home/tim/miniconda3/envs/flux/lib/python3.11/site-packages/accelerate/commands/accelerate_cli.py", line 48, in main
    args.func(args)
  File "/home/tim/miniconda3/envs/flux/lib/python3.11/site-packages/accelerate/commands/launch.py", line 1159, in launch_command
    multi_gpu_launcher(args)
  File "/home/tim/miniconda3/envs/flux/lib/python3.11/site-packages/accelerate/commands/launch.py", line 793, in multi_gpu_launcher
    distrib_run.run(args)
  File "/home/tim/miniconda3/envs/flux/lib/python3.11/site-packages/torch/distributed/run.py", line 910, in run
    elastic_launch(
  File "/home/tim/miniconda3/envs/flux/lib/python3.11/site-packages/torch/distributed/launcher/api.py", line 138, in __call__
    return launch_agent(self._config, self._entrypoint, list(args))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tim/miniconda3/envs/flux/lib/python3.11/site-packages/torch/distributed/launcher/api.py", line 269, in launch_agent
    raise ChildFailedError(
torch.distributed.elastic.multiprocessing.errors.ChildFailedError: 
========================================================
train_dreambooth_flux.py FAILED
--------------------------------------------------------
Failures:
  <NO_OTHER_FAILURES>
--------------------------------------------------------
Root Cause (first observed failure):
[0]:
  time      : 2024-12-17_15:48:37
  host      : chariot
  rank      : 1 (local_rank: 1)
  exitcode  : -9 (pid: 1629301)
  error_file: <N/A>
  traceback : Signal 9 (SIGKILL) received by PID 1629301
========================================================

I'm using two GeForce RTX4090 (24GB) so in theory it should work with accelerate?

I'm following the script as well, so seems like something is not right?
Do you need more information about the environment for better guidance?
It seems that the model is somehow loaded only into the cpu, is there an option to circumvent that?
Could it be the accelerate config (default)?

Thanks in advance!

[Becomebright]

Any update?

[liiiiiiiiil]

using deepspeed，refer to https://github.com/huggingface/diffusers/blob/main/examples/flux-control/train_control_flux.py

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.