[Tracker] fix training resuming problem when using FP16 in the examples

[sayakpaul]

We have been a lot of issues on this topic:

 File "/content/train_text_to_image_lora_sdxl.py", line 1261, in <module>
 main(args)
 File "/content/train_text_to_image_lora_sdxl.py", line 1077, in main
 accelerator.clip_grad_norm_(params_to_optimize, args.max_grad_norm)
 File "/usr/local/lib/python3.10/dist-packages/accelerate/accelerator.py", line 2040, in clip_grad_norm_
 self.unscale_gradients()
 File "/usr/local/lib/python3.10/dist-packages/accelerate/accelerator.py", line 2003, in unscale_gradients
 self.scaler.unscale_(opt)
 File "/usr/local/lib/python3.10/dist-packages/torch/cuda/amp/grad_scaler.py", line 307, in unscale_
 optimizer_state["found_inf_per_device"] = self._unscale_grads_(
 File "/usr/local/lib/python3.10/dist-packages/torch/cuda/amp/grad_scaler.py", line 229, in _unscale_grads_
 raise ValueError("Attempting to unscale FP16 gradients.")

• ValueError: Attempting to unscale FP16 gradients. #6442
• ValueError: Attempting to unscale FP16 gradients #6098
• SDXL dreambooth can't be resumed from a checkpoint at fp16 training #5004

#6514 introduces a fix for the SDXL DreamBooth LoRA training script. #6553 is a follow-up that cleans things up a bit. However, it would be nice to extend them for the other scripts too. This issue tracks the integration:

• DreamBooth LoRA SD
• Text-to-image LoRA SDXL
• SDXL Consistency Distillation
• Advanced LoRA trainer (cc: @linoytsaban)
• Advanced LoRA trainer SD v1.5 (cc: @linoytsaban)

Feel free to claim any of these and submit PRs. Do tag me in those PRs and focus on fixing only one script at a time. I know that @linoytsaban is already working on the last one.

While submitting the PRs, please also provide example commands as I did in #6514 (comment) so as to quickly test the validity.

[a-r-r-o-w]

Hi Sayak, I'll take up DreamBooth LoRA SD. Opening a PR shortly, thanks.

[steverhoades]

So just an FYI i did a quick port of those changes to the Advanced LoRa trainer. I received the following error after resuming training from a checkpoint and running inference on the generated model.

Loading adapter weights from state_dict led to unexpected keys not found in the model: ['down_blocks.1.attentions.0.transformer_blocks.0.attn1.to_q.lora_A_1.default_0.weight', 'down_blocks.1.attentions.0.transformer_blocks.0.attn1.to_q.lora_B_1.default_0.weight',
....

Update:
I can see in the logs that once training finishes the same error is raised.

[steverhoades]

Tested training a model with train_dreambooth_lora_sdxl.py script and the following error was returned after saving a checkpoint and trying to resume from the checkpoint.

#!/usr/bin/env bash
!accelerate launch scripts/train_dreambooth_lora_sdxl.py \
 --report_to="wandb" \
 --pretrained_model_name_or_path="stabilityai/stable-diffusion-xl-base-1.0" \
 --pretrained_vae_model_name_or_path="madebyollin/sdxl-vae-fp16-fix" \
 --dataset_name="./training_images" \
 --output_dir="father_lora_v22" \
 --cache_dir="./dataset_cache_dir" \
 --caption_column="prompt" \
 --mixed_precision="fp16" \
 --instance_prompt="a photo of Brian de palma" \
 --resolution=1024 \
 --train_batch_size=1 \
 --gradient_accumulation_steps=4 \
 --gradient_checkpointing \
 --snr_gamma=5.0 \
 --lr_scheduler="cosine_with_restarts" \
 --lr_warmup_steps=0 \
 --repeats=10 \
 --max_train_steps=20 \
 --checkpointing_steps=10 \
 --validation_prompt="a photo of Brian de palma" \
 --validation_epochs=1 \
 --with_prior_preservation \
 --class_data_dir="./prior_preservation-man-v2" \
 --num_class_images=110 \
 --class_prompt="a photo of a man" \
 --rank=32 \
 --optimizer="prodigy" \
 --prodigy_safeguard_warmup=True \
 --prodigy_use_bias_correction=True \
 --adam_beta1=0.9 \
 --adam_beta2=0.99 \
 --adam_weight_decay=0.01 \
 --train_text_encoder \
 --learning_rate=1 \
 --text_encoder_lr=1 \
 --resume_from_checkpoint="checkpoint-10" \
 --seed="0"

resulted in:

Resuming from checkpoint checkpoint-10
01/12/2024 16:42:50 - INFO - accelerate.accelerator - Loading states from father_lora_v22/checkpoint-10
Traceback (most recent call last):
 File "/home/steve/projects/training/scripts/train_dreambooth_lora_sdxl.py", line 1762, in <module>
 main(args)
 File "/home/steve/projects/training/scripts/train_dreambooth_lora_sdxl.py", line 1383, in main
 accelerator.load_state(os.path.join(args.output_dir, path))
 File "/home/steve/.local/lib/python3.10/site-packages/accelerate/accelerator.py", line 2849, in load_state
 hook(models, input_dir)
 File "/home/steve/projects/training/scripts/train_dreambooth_lora_sdxl.py", line 1078, in load_model_hook
 _set_state_dict_into_text_encoder(
 File "/home/steve/.local/lib/python3.10/site-packages/diffusers/training_utils.py", line 140, in _set_state_dict_into_text_encoder
 set_peft_model_state_dict(text_encoder, text_encoder_state_dict, adapter_name="default")
 File "/home/steve/.local/lib/python3.10/site-packages/peft/utils/save_and_load.py", line 201, in set_peft_model_state_dict
 load_result = model.load_state_dict(peft_model_state_dict, strict=False)
 File "/home/steve/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 2152, in load_state_dict
 raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
RuntimeError: Error(s) in loading state_dict for CLIPTextModel:
size mismatch for text_model.encoder.layers.0.self_attn.k_proj.lora_A.default.weight: copying a param with shape torch.Size([32, 1280]) from checkpoint, the shape in current model is torch.Size([32, 768]).
	size mismatch for text_model.encoder.layers.0.self_attn.k_proj.lora_B.default.weight: copying a param with shape torch.Size([1280, 32]) from checkpoint, the shape in current model is torch.Size([768, 32]).
 ...

[sayakpaul]

@linoytsaban ^

[steverhoades]

I was able to get a working example running, still learning this code base but gave a crack at a PR #6566.

I tested both train_text_encoder and train_text_encoder_ti which completed successfully.

!accelerate launch scripts/train_dreambooth_lora_sdxl_advanced_orig.py \
 --report_to="wandb" \
 --pretrained_model_name_or_path="stabilityai/stable-diffusion-xl-base-1.0" \
 --pretrained_vae_model_name_or_path="madebyollin/sdxl-vae-fp16-fix" \
 --dataset_name="./training_set" \
 --output_dir="father_lora_v21" \
 --cache_dir="./dataset_cache_dir" \
 --caption_column="prompt" \
 --mixed_precision="fp16" \
 --instance_prompt="a photo of Brian de palma" \
 --resolution=1024 \
 --train_batch_size=1 \
 --gradient_accumulation_steps=4 \
 --gradient_checkpointing \
 --snr_gamma=5.0 \
 --lr_scheduler="cosine_with_restarts" \
 --lr_warmup_steps=0 \
 --repeats=10 \
 --max_train_steps=20 \
 --checkpointing_steps=10 \
 --validation_prompt="a photo of Brian de palma in a suit, looking at the camera" \
 --validation_epochs=1 \
 --with_prior_preservation \
 --class_data_dir="./prior_preservation-man-v2" \
 --num_class_images=110 \
 --class_prompt="a photo of a man" \
 --rank=32 \
 --optimizer="prodigy" \
 --prodigy_safeguard_warmup=True \
 --prodigy_use_bias_correction=True \
 --adam_beta1=0.9 \
 --adam_beta2=0.99 \
 --adam_weight_decay=0.01 \
 --train_text_encoder \
 --learning_rate=1 \
 --text_encoder_lr=1 \
 --resume_from_checkpoint="checkpoint-10" \
 --seed="0"

!accelerate launch scripts/train_dreambooth_lora_sdxl_advanced.py \
 --report_to="wandb" \
 --pretrained_model_name_or_path="stabilityai/stable-diffusion-xl-base-1.0" \
 --dataset_name="./training_set_ti" \
 --output_dir="father_lora_v1-test" \
 --cache_dir="./dataset_cache_dir" \
 --caption_column="prompt" \
 --mixed_precision="fp16" \
 --instance_prompt="a photo of TOK man" \
 --resolution=1024 \
 --train_batch_size=1 \
 --gradient_accumulation_steps=4 \
 --gradient_checkpointing \
 --snr_gamma=5.0 \
 --lr_scheduler="constant" \
 --lr_warmup_steps=0 \
 --max_train_steps=20 \
 --checkpointing_steps=10 \
 --checkpoints_total_limit=10 \
 --validation_prompt="a photo of TOK man in a suit, looking directly at the camera" \
 --validation_epochs=1 \
 --with_prior_preservation \
 --class_data_dir="./prior_preservation-man-v2" \
 --num_class_images=100 \
 --class_prompt="a photo of man" \
 --rank=32 \
 --optimizer="prodigy" \
 --prodigy_safeguard_warmup=True \
 --prodigy_use_bias_correction=True \
 --adam_beta1=0.9 \
 --adam_beta2=0.99 \
 --adam_weight_decay=0.01 \
 --learning_rate=1 \
 --text_encoder_lr=1 \
 --train_text_encoder_ti \
 --train_text_encoder_frac=0.5 \
 --resume_from_checkpoint="checkpoint-10" \
 --seed="0"

[pedrogengo]

Hey hey! I will take Text-to-image LoRA SD fix. Should open a PR today or tomorrow!

[sriagastya]

Hello HuggingFace Team, Can I work on SDXL Consistency Distillation?