train_dreambooth_lora_flux crash on batch size greater than 1

[PluginBOXone]

Describe the bug

As soon as launching the train_dreambooth_lora_flux.py (with accelerate) and train_batch_size greater than 1, it takes 40-50sec to load and then crashes with an error "RuntimeError: mat1 and mat2 shapes cannot be multiplied (1x1536 and 768x3072)"

File "/workspace/./train_dreambooth_lora_flux.py", line 1935, in main(args) File "/workspace/./train_dreambooth_lora_flux.py", line 1729, in main model_pred = transformer( File "/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl return self._call_impl(*args, **kwargs) File "/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py", line 1562, in _call_impl return forward_call(*args, **kwargs) File "/usr/local/lib/python3.11/dist-packages/accelerate/utils/operations.py", line 819, in forward return model_forward(*args, **kwargs) File "/usr/local/lib/python3.11/dist-packages/torch/amp/autocast_mode.py", line 43, in decorate_autocast return func(*args, **kwargs) File "/usr/local/lib/python3.11/dist-packages/diffusers/models/transformers/transformer_flux.py", line 493, in forward else self.time_text_embed(timestep, guidance, pooled_projections) File "/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl return self._call_impl(*args, **kwargs) File "/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py", line 1562, in _call_impl return forward_call(*args, **kwargs) File "/usr/local/lib/python3.11/dist-packages/diffusers/models/embeddings.py", line 1629, in forward pooled_projections = self.text_embedder(pooled_projection) File "/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl return self._call_impl(*args, **kwargs) File "/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py", line 1562, in _call_impl return forward_call(*args, **kwargs) File "/usr/local/lib/python3.11/dist-packages/diffusers/models/embeddings.py", line 2222, in forward hidden_states = self.linear_1(caption) File "/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl return self._call_impl(*args, **kwargs) File "/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py", line 1562, in _call_impl return forward_call(*args, **kwargs) File "/usr/local/lib/python3.11/dist-packages/torch/nn/modules/linear.py", line 117, in forward return F.linear(input, self.weight, self.bias)

Reproduction

FLUX1 Dev used

accelerate launch --config_file /workspace/accelerate.yaml ./train_dreambooth_lora_flux.py --pretrained_model_name_or_path /workspace/model/realflux1 --output_dir /workspace/output_model/577cda6e-628d-4dc0-9ab4-cbcf448b730e-e1 --resolution 1024 --learning_rate 5e-4 --mixed_precision bf16 --instance_data_dir /workspace/job_files/577cda6e-628d-4dc0-9ab4-cbcf448b730e-e1/clean_data --instance_prompt WIXBSAHA black car --lr_warmup_steps 0 --gradient_accumulation_steps 1 --lr_scheduler cosine --train_batch_size 2 --max_train_steps 500 --use_8bit_adam --gradient_checkpointing --checkpointing_steps 500 --num_train_epochs 10 --checkpoints_total_limit 1 --train_text_encoder --rank 16 --guidance_scale 1 --cache_latents

Logs

Num processes: 1
Process index: 0
Local process index: 0
Device: cuda
Mixed precision type: bf16
You set `add_prefix_space`. The tokenizer needs to be converted from the slow tokenizers
You are using a model of type clip_text_model to instantiate a model of type . This is not supported for all configurations of models and can yield errors.
You are using a model of type t5 to instantiate a model of type . This is not supported for all configurations of models and can yield errors.
{'use_beta_sigmas', 'shift_terminal', 'invert_sigmas', 'use_exponential_sigmas', 'use_karras_sigmas'} was not found in config. Values will be initialized to default values.
{'out_channels', 'axes_dims_rope'} was not found in config. Values will be initialized to default values.

System Info

RunPod
NVIDIA A100 80GB PCIe CUDA Version: 12.4
accelerate 0.33.0
bitsandbytes 0.43.1
datasets 2.19.2
diffusers 0.33.0.dev0
Jinja2 3.1.6
peft 0.14.0
safetensors 0.5.3
sentencepiece 0.2.0
tensorboard 2.19.0
tokenizers 0.21.0
torch 2.4.1+cu124
torchaudio 2.4.1+cu124
torchvision 0.19.1+cu124
transformers 4.49.0

Who can help?

@sayakpaul

[sayakpaul]

Can you try to reproduce it with an example command from here? Otherwise, it won't be possible for us to look into it further.

Cc: @linoytsaban

[PluginBOXone]

I tried following command 1:1 expect train_batch_size on 2 and now it says "CUDA out of memory. Tried to allocate 108.00 MiB. GPU 0 has a total capacity of 79.15 GiB of which 57.25 MiB is free. Process 924987 has 79.08 GiB memory in use. Of the allocated memory 78.20 GiB is allocated by PyTorch, and 395.28 MiB is reserved by PyTorch but unallocated."

accelerate launch train_dreambooth_lora_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=512
--train_batch_size=2
--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

Is FLUX1 Dev really taking up more than 80GB of VRAM at batch size 2???

[sayakpaul]

Sorry for the late reply.

I used the following on a single A100 and was able to run successfully:

CUDA_VISIBLE_DEVICES=0 accelerate launch train_dreambooth_lora_flux.py \
  --pretrained_model_name_or_path="black-forest-labs/FLUX.1-dev"\
  --instance_data_dir="./dog"\
  --output_dir="trained-flux-lora"\
  --mixed_precision="bf16"\
  --instance_prompt="a photo of sks dog"\
  --mixed_precision="bf16" \
  --resolution=1024 \
  --train_batch_size=1\
  --guidance_scale=1\
  --gradient_accumulation_steps=4\
  --optimizer="adamw"\
  --use_8bit_adam\
  --learning_rate=1e-4\
  --lr_scheduler="constant"\
  --lr_warmup_steps=0\
  --max_train_steps=100\
  --seed="0"

My diffusers-cli env:

- 🤗 Diffusers version: 0.33.0.dev0
- Platform: Linux-5.4.0-166-generic-x86_64-with-glibc2.31
- Running on Google Colab?: No
- Python version: 3.10.12
- PyTorch version (GPU?): 2.6.0+cu124 (True)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Huggingface_hub version: 0.27.0
- Transformers version: 4.49.0
- Accelerate version: 1.2.1
- PEFT version: 0.14.0
- Bitsandbytes version: 0.45.3
- Safetensors version: 0.4.4
- xFormers version: 0.0.28.post3
- Accelerator: NVIDIA A100-SXM4-80GB, 81920 MiB
NVIDIA A100-SXM4-80GB, 81920 MiB
NVIDIA A100-SXM4-80GB, 81920 MiB
NVIDIA DGX Display, 4096 MiB
NVIDIA A100-SXM4-80GB, 81920 MiB
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: No

[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.

[a-r-r-o-w]

Marking as closed due to inactivity and because @sayakpaul could not repro