[HiDream LoRA] optimizations + small updates

[linoytsaban]

some memory optimizations:

• add pre-computation of prompt embeddings when custom prompts are used as well
• add pre-computation of validation prompt as well
• add --skip_final_inference - to allow to run with validation, but skip the final loading of the pipeline with the lora weights to reduce memory reqs

other changes:

• update default trained layers
• save model card even if model is not pushed to hub
• remove scheduler initialization from code example - not necessary anymore (it's now if the base model's config)

todo:

• update readme with better defaults

Yarn Art LoRA

training config

import os
os.environ['MODEL_NAME'] = "HiDream-ai/HiDream-I1-Full"
os.environ['DATASET_NAME'] ="Norod78/Yarn-art-style"
os.environ['OUTPUT_DIR'] = "hidream-yarn-art-lora-v2-trainer"

!accelerate launch train_dreambooth_lora_hidream.py \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --dataset_name=$DATASET_NAME \
  --output_dir=$OUTPUT_DIR \
  --mixed_precision="bf16" \
  --lora_layers="to_k,to_q,to_v,to_out"\
  --instance_prompt="a dog, yarn art style" \
  --validation_prompt="yoda, yarn art style" \
  --caption_column="text" \
  --resolution=1024 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=1 \
  --use_8bit_adam\
  --rank=16 \
  --learning_rate=1e-4 \
  --report_to="wandb" \
  --lr_scheduler="constant_with_warmup" \
  --lr_warmup_steps=200 \
  --max_train_steps=1000 \
  --validation_epochs=25 \
  --seed="0" \
  --push_to_hub

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[linoytsaban]

@sayakpaul now that I'm thinking about it, even when validation is enabled, since we're not optimizing the text encoders, can't we just pre-encode the validation prompt embeddings as well? and then we don't need to keep or load text encoders for validation at all and simply pass the embeddings to log_validation

[sayakpaul]

@sayakpaul now that I'm thinking about it, even when validation is enabled, since we're not optimizing the text encoders, can't we just pre-encode the validation prompt embeddings as well? and then we don't need to keep or load text encoders for validation at all and simply pass the embeddings to log_validation

We should. Let's do that.

[sayakpaul]

Thanks for the optims. LMK if the comments make sense or if anything is unclear.

[linoytsaban]

@bot /style

[github-actions]

Style fixes have been applied. View the workflow run here.

[sayakpaul]

Thank you! Left some more comments. I think we can merge this today. Also would be good to see if we test this on a 40GB GPU.

[linoytsaban]

Thanks @sayakpaul!

I added prints using your snippet (here) - right before caching & and pre-computation of prompt embeddings and straight after deleting them.
freeing the memory seems to work as expected getting us to ~33GB, but we hit ~60GB prior when we move the text encoding pipeline to gpu (this is with resolution==1024)
with offloading -

=== CUDA Memory Stats before caching ===
Current allocated: 32.24 GB
Max allocated: 32.24 GB
Current reserved: 58.29 GB
Max reserved: 58.29 GB

=== CUDA Memory Stats after caching ===
Current allocated: 58.41 GB
Max allocated: 58.41 GB
Current reserved: 60.84 GB
Max reserved: 60.84 GB

=== CUDA Memory Stats after freeing ===
Current allocated: 32.87 GB
Max allocated: 32.87 GB
Current reserved: 33.82 GB
Max reserved: 33.82 GB

without offloading & caching -

=== CUDA Memory Stats before caching ===
Current allocated: 57.77 GB
Max allocated: 57.77 GB
Current reserved: 58.50 GB
Max reserved: 58.50 GB

=== CUDA Memory Stats after caching ===
Current allocated: 58.35 GB
Max allocated: 58.35 GB
Current reserved: 59.11 GB
Max reserved: 59.11 GB

=== CUDA Memory Stats after freeing ===
Current allocated: 32.98 GB
Max allocated: 32.98 GB
Current reserved: 33.39 GB
Max reserved: 33.39 GB

[linoytsaban]

@bot /style

[github-actions]

Style fixes have been applied. View the workflow run here.

[linoytsaban]

@bot /style

[github-actions]

Style fixes have been applied. View the workflow run here.

[sayakpaul]

Looks great! Thanks Linoy!

[sayakpaul]

Perhaps we can add a comment explaining that including to_out will target all the expert layers.

[linoytsaban]

@bot /style