[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