CudaOutOfMemory on Flux Lora Training #9156
Description
Describe the bug
I tried to train the flux-dev model with Lora on A100 40GB. But it raises the CudaOutOfMemory exception.
Reproduction
# Accelerate command
export MODEL_NAME="black-forest-labs/FLUX.1-dev"
export INSTANCE_DIR="woman"
export OUTPUT_DIR="trained-flux-lora-woman"
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="fp16" \
--instance_prompt="ohwx woman" \
--resolution=512 \
--train_batch_size=1 \
--gradient_accumulation_steps=4 \
--learning_rate=1e-5 \
--lr_scheduler="constant" \
--lr_warmup_steps=0 \
--max_train_steps=500 \
--validation_prompt="professional photography of ohwx woman" \
--validation_epochs=25 \
--seed="0" \
--use_8bit_adam
# Accelerate Config
compute_environment: LOCAL_MACHINE
debug: false
distributed_type: 'NO'
downcast_bf16: 'no'
enable_cpu_affinity: true
gpu_ids: all
machine_rank: 0
main_training_function: main
mixed_precision: fp16
num_machines: 1
num_processes: 1
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false
Logs
[Previous line repeated 4 more times]
File "/home/muhammed_pektas/anaconda3/envs/hflora/lib/python3.12/site-packages/torch/nn/modules/module.py", line 805, in _apply
param_applied = fn(param)
^^^^^^^^^
File "/home/muhammed_pektas/anaconda3/envs/hflora/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1160, in convert
return t.to(
^^^^^
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 64.00 MiB. GPU 0 has a total capacity of 39.39 GiB of which 14.31 MiB is free. Including non-PyTorch memory, this process has 39.37 GiB memory in use. Of the allocated memory 38.79 GiB is allocated by PyTorch, and 73.66 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.System Info
Pip
req_flux.txt
Hardware
NVIDIA A100-SXM4-40GB
Who can help?
No response