Dreambooth Flux with Lora: OOM

[massyzs]

Describe the bug

It is pretty wired that even after I set lora rank to only 2, training the dreambooth_flux with lora still returns OOM. My GPU is three L40s (44G-48G). And data is only one image.

If this is still OOM, I'm wondering why should lora version exist? The next GPU memory capbility is 80G which can train it without lora.

Reproduction

NA

Logs

No response

System Info

L40s

Who can help?

No response

[hlky]

cc @linoytsaban @sayakpaul

[sayakpaul]

What's your training command?

If this is still OOM, I'm wondering why should lora version exist?

Maybe we could tone it down a little bit?

[massyzs]

This is the training script for dreambooth_flux_lora, I modify the lora rank from default 4 to 2 to save memory and even try to add "8-bit optimizer" which seems to be of no use.

accelerate launch train_dreambooth_flux_lora.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 bear" \ --resolution=512 \ --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 bear in forest" \ --validation_epochs=25 \ --seed="0" \

May I know about the possible reason or how to solve this problem? Many thanks.

[sayakpaul]

Would suggest adding the following CLI args to your command:

+ --gradient_checkpointing \
+ --cache_latents \
+ --use_8bit_adam \

Cc: @linoytsaban

[massyzs]

Thanks for your reply, now the script goes like this:

CUDA_VISIBLE_DEVICES=0 accelerate launch train_dreambooth_flux_lora.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 bear" \ --resolution=512 \ --train_batch_size=1 \ --guidance_scale=1 \ --gradient_accumulation_steps=1 \ --learning_rate=1. \ --report_to="wandb" \ --lr_scheduler="constant" \ --lr_warmup_steps=0 \ --max_train_steps=500 \ --validation_prompt="A photo of sks bear in forest" \ --validation_epochs=25 \ --seed="0" \ --gradient_checkpointing \ --cache_latents \ --use_8bit_adam \

the output is like:

It is still OOM which is pretty wired. Do you have any idea on this?

[sayakpaul]

Will let @linoytsaban further comment. If this is a blocker and cannot wait till Linoy replies check out https://github.com/ostris/ai-toolkit/.

[massyzs]

Thanks for your help.
I guess it is because the prompt encoding steps I added in validation step:
`

 prompt_embeds, pooled_prompt_embeds, text_ids = pipeline.encode_prompt(
          args.validation_prompt, prompt_2=args.validation_prompt
      )

with autocast_ctx:
    images = [pipeline(
                prompt_embeds=prompt_embeds,
                pooled_prompt_embeds=pooled_prompt_embeds,
               
                num_inference_steps=28,
                controlnet_conditioning_scale=0.7,
                guidance_scale=3.5,
                generator=generator).images[0] for _ in range(args.num_validation_images)]

`
I find this from other github issues (I guess you remember this). When goes to validation, it even temporarily goes to 90G. And so far I cannot remove this for "T5 does not work"

[sayakpaul]

Does it fail here?

 prompt_embeds, pooled_prompt_embeds, text_ids = pipeline.encode_prompt(
          args.validation_prompt, prompt_2=args.validation_prompt
      )

A couple of things you could try further (some of them will require changes to the code):

• Disable validation inference
• After encode_prompt() is called delete, the text encoders (del pipe.text_encoder, pipe.text_encoder_2)
• Instead of placing the pipeline on accelerate.device call enable_model_cpu_offload() on pipe.

[linoytsaban]

I agree validation is likely the main culprit behind this, if I may ask, what's the motivation behind the additional prompt encoding?
either way, even w/o the additional encoding and as @sayakpaul mentioned, disabling validation would likely release a significant chunk of memory - but I'd try @sayakpaul suggestions first. For context I trained Flux LoRAs of rank 16-32 with text encoder training and prodigy optimizer on an A100 w/o OOMing using this script - but I indeed disabled validation

[massyzs]

Hi, thanks for replying. The motivation is because T5 does not support the auto casting. This solution is actually given by sayakpaul. Now, the 8-bit adam and gradient checkpointing (suggest check is there any problem) seems to be of no use. While it indeed can work if without validation. 发自我的iPhone

…

------------------ Original ------------------ From: Linoy Tsaban ***@***.***> Date: Sun, Dec 8, 2024 6:10 PM To: huggingface/diffusers ***@***.***> Cc: 曹潇 ***@***.***>, Author ***@***.***> Subject: Re: [huggingface/diffusers] Dreambooth Flux with Lora: OOM (Issue#10145) I agree validation is likely the main culprit behind this, if I may ask, what's the motivation behind the additional prompt encoding? either way, even w/o the additional encoding and as @sayakpaul mentioned, disabling validation would likely release a significant chunk of memory - but I'd try @sayakpaul suggestions first. For context I trained Flux LoRAs of rank 16-32 with text encoder training and prodigy optimizer on an A100 w/o OOMing using this script - but I indeed disabled validation — Reply to this email directly, view it on GitHub, or unsubscribe. You are receiving this because you authored the thread.Message ID: ***@***.***>

[massyzs]

I agree validation is likely the main culprit behind this, if I may ask, what's the motivation behind the additional prompt encoding? either way, even w/o the additional encoding and as @sayakpaul mentioned, disabling validation would likely release a significant chunk of memory - but I'd try @sayakpaul suggestions first. For context I trained Flux LoRAs of rank 16-32 with text encoder training and prodigy optimizer on an A100 w/o OOMing using this script - but I indeed disabled validation

Hi, I think the validation code and test code (inference) are really in not useable state. I really suggest you check and test before releasing the code, especially if you are on behalf of huggingface diffusers.

[sayakpaul]

Hi, I think the validation code and test code (inference) are really in not useable state.

When is it not useable?