train with train_text_to_image_sdxl.py CUDA out of memory

[MMY1994]

This is my shell script

export MODEL_NAME="./model_base/sd_xl_base"
export VAE_NAME="./model_base/sdxl-vae-fp16-fix"

accelerate launch
--config_file /root/.cache/huggingface/accelerate/default_config.yaml
--num_cpu_threads_per_process=8
./diffusers/examples/text_to_image/train_text_to_image_sdxl_mmy.py
--pretrained_model_name_or_path=$MODEL_NAME
--pretrained_vae_model_name_or_path=$VAE_NAME
--train_data_dir="./fineturn/fineturn.py"
--output_dir="./sdxl-dreamshape-fineturn-model/"
--enable_xformers_memory_efficient_attention
--resolution=1024 --center_crop --random_flip
--proportion_empty_prompts=0.2
--train_batch_size=1
--mixed_precision="fp16"
--gradient_accumulation_steps=4 --gradient_checkpointing
--max_train_steps=5000
--use_8bit_adam
--learning_rate=1e-07
--lr_scheduler="constant" --lr_warmup_steps=0
--validation_prompt="a young man climbing the ladder, high quality" --validation_epochs 1
--checkpointing_steps=200

I Use A800 40GB and --enable_xformers_memory_efficient_attention, --mixed_precision="fp16".

Even when I use 512 resolution and batchsize=1, this error still happen. i don't think it's normal

[tolgacangoz]

Could you share the output of diffusers-cli env command? I see that you added a suffix to the file name and main(args) is at 1333. line, is there any change in the file in this discussion's context? I wonder if --num_cpu_threads_per_process=8 tries to allocate too much data? Also, what happens when you remove --enable_xformers_memory_efficient_attention if your PyTorch>=2.0.

[tolgacangoz]

1. PagedAdamW8bit uses system RAM when OOM seems to occur in GPU, so it is normal to use all GPU vRAM -it uses all GPU vRAM + less system RAM with a batch size of 1, and uses all GPU vRAM + more system RAM with a batch size of 8, AFAIU. For the further batch sizes, the deepspeed library could be researched.
2. I couldn't reproduce this with the setup below. Maybe, your torchvision and/or datasets are problematic; upgrade them?

~25 GB system RAM used at peak, 24 GB vRAM NVIDIA L4.
Command:

accelerate launch
--num_cpu_threads_per_process=1
./diffusers/examples/text_to_image/train_text_to_image_sdxl.py
--pretrained_model_name_or_path="stabilityai/stable-diffusion-xl-base-1.0"
--pretrained_vae_model_name_or_path="madebyollin/sdxl-vae-fp16-fix"
--train_data_dir="YaYaB/onepiece-blip-captions"
--output_dir="./sdxl-dreamshape-fineturn-model/"
--resolution=1024 --center_crop --random_flip
--proportion_empty_prompts=0.2
--train_batch_size=1
--mixed_precision="fp16"
--gradient_checkpointing
--max_train_steps=10
--use_8bit_adam
--learning_rate=1e-07
--lr_scheduler="constant" --lr_warmup_steps=0
--validation_prompt="a young man climbing the ladder, high quality" --validation_epochs 1
--checkpointing_steps=200

with this modification in the examples/text_to_image/train_text_to_image_sdxl.py at the 790. line:

-        optimizer_class = bnb.optim.AdamW8bit
+        optimizer_class = bnb.optim.PagedAdamW8bit

accelerate env:

- `Accelerate` default config:
        - compute_environment: LOCAL_MACHINE
        - distributed_type: NO
        - mixed_precision: fp16
        - use_cpu: False
        - debug: False
        - num_processes: 1
        - machine_rank: 0
        - num_machines: 1
        - gpu_ids: all
        - rdzv_backend: static
        - same_network: True
        - main_training_function: main
        - enable_cpu_affinity: False
        - downcast_bf16: no
        - tpu_use_cluster: False
        - tpu_use_sudo: False
        - tpu_env: []

diffusers-cli env:

- `diffusers` version: 0.28.0.dev0
- Platform: Linux-5.15.0-1058-aws-x86_64-with-glibc2.31
- Python version: 3.10.10
- PyTorch version (GPU?): 2.4.0.dev20240507+cu121 (True)
- Huggingface_hub version: 0.23.0
- Transformers version: 4.40.2
- Accelerate version: 0.30.0
- xFormers version: not installed
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

python -c "import bitsandbytes as bnb, torchvision as tv, datasets as ds; print('bnb:', bnb.__version__); print('torchvision:', tv.__version__); print('datasets:', ds.__version__);"
bnb: 0.43.1
torchvision: 0.19.0.dev20240507+cu121
datasets: 2.19.1

[MMY1994]

Thank you, I will try it

[MMY1994]

After try it, GPU memery problem seems to be related to the torch version(my torch version is 2.0.1)。accelerate config's distributed_type set to "NO" is norml( single GPU about 25GB), but set to "MULTI_GPU"is about 38GB. it reported the Info "torch.nn.parallel.distributed - Reducer buckets have been rebuilt in this iteration".
I will try upgrade them

[MMY1994]

I'm sorry to interrupt, but have you tried single machine multi GPU and distribution, like this

accelerate env:

• Accelerate default config:
  - compute_environment: LOCAL_MACHINE
  - distributed_type: MULTI_GPU
  - mixed_precision: fp16
  - use_cpu: False
  - debug: False
  - num_processes: 8
  - machine_rank: 0
  - num_machines: 1
  - gpu_ids: all
  - rdzv_backend: static
  - same_network: True
  - main_training_function: main
  - enable_cpu_affinity: False
  - downcast_bf16: no
  - tpu_use_cluster: False
  - tpu_use_sudo: False
  - tpu_env: []

According to the results of my experiment, single GPU is normal, but multi-GPUs still takes up too much GPU memory

[tolgacangoz]

AFAIU, PagedAdamW8bit uses as much GPU vRAM as possible. If there seems to be a possibility of OOM, then at that point, it starts to utilize CPU RAM. If multi-GPUs have more vRAM than single GPU then multi-GPUs will use more vRAM according to training needs.
Actually, my knowledge about distributed training is very limited. It might be more proper to open a new discussion about that.