Getting CUDA out of memory error even with Colab A100 high RAM

### Describe the bug

I am trying to fine-tune Flux.1-dev with Lora on the Google Colab A100 runtime environment. It has 80 GB system RAM and 40 GB VRAM. I followed the recommended steps from this [link](https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/README_flux.md). I am still getting "CUDA out of memory" error. I saw that related old issues were closed, but the bug seems to be still present.

### Reproduction

Here is the last version that I tried. It is with AdamQ. I also tried using Prodigy optimizer and got the same error.

!accelerate launch -q train_dreambooth_lora_flux.py \
--pretrained_model_name_or_path="black-forest-labs/FLUX.1-dev" \
--instance_data_dir="train_photos" \
--output_dir="trained-flux-lora" \
--mixed_precision="bf16" \
--instance_prompt="a photo of X" \
--resolution=512 \
--rank=1 \
--train_batch_size=1 \
--guidance_scale=1 \
--gradient_accumulation_steps=4 \
--optimizer="AdamW" \
--learning_rate=1. \
--lr_scheduler="constant" \
--lr_warmup_steps=0 \
--max_train_steps=500 \
--validation_prompt="A photo of X biking" \
--validation_epochs=25 \
--seed="0" \
--lora_layers="attn.to_k,attn.to_q,attn.to_v,attn.to_out.0" \
--use_8bit_adam

### Logs

```shell
All the weights of FluxTransformer2DModel were initialized from the model checkpoint at black-forest-labs/FLUX.1-dev.
If your task is similar to the task the model of the checkpoint was trained on, you can already use FluxTransformer2DModel for predictions without further training.
03/09/2025 15:45:57 - INFO - __main__ - ***** Running training *****
03/09/2025 15:45:57 - INFO - __main__ - Num examples = 5
03/09/2025 15:45:57 - INFO - __main__ - Num batches each epoch = 5
03/09/2025 15:45:57 - INFO - __main__ - Num Epochs = 250
03/09/2025 15:45:57 - INFO - __main__ - Instantaneous batch size per device = 1
03/09/2025 15:45:57 - INFO - __main__ - Total train batch size (w. parallel, distributed & accumulation) = 4
03/09/2025 15:45:57 - INFO - __main__ - Gradient Accumulation steps = 4
03/09/2025 15:45:57 - INFO - __main__ - Total optimization steps = 500
Steps: 0% 0/500 [00:00
main(args)
File "/content/train_dreambooth_lora_flux.py", line 1720, in main
model_pred = transformer(
^^^^^^^^^^^^
File "/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py", line 1747, 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/accelerate/utils/operations.py", line 807, in __call__
return convert_to_fp32(self.model_forward(*args, **kwargs))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.11/dist-packages/torch/amp/autocast_mode.py", line 44, in decorate_autocast
return func(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.11/dist-packages/diffusers/models/transformers/transformer_flux.py", line 523, in forward
hidden_states = block(
^^^^^^
File "/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py", line 1747, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.11/dist-packages/diffusers/models/transformers/transformer_flux.py", line 96, in forward
hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 90.00 MiB. GPU 0 has a total capacity of 39.56 GiB of which 46.88 MiB is free. Process 68220 has 39.50 GiB memory in use. Of the allocated memory 38.84 GiB is allocated by PyTorch, and 172.60 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation. See documentation for Memory Management (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)
Steps: 0% 0/500 [00:01

Reproduce the command in a Google Colab A100 environment, then inspect train_dreambooth_lora_flux.py around line 1720 and diffusers/models/transformers/transformer_flux.py around lines 523 and 96. Compare the reported GPU allocation and identify why the recommended configuration exceeds 40 GB; done means training completes without the CUDA out-of-memory failure.

Escrito por el modelo de indexación a partir del texto del issue.