[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!

[asrimanth]

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