Skip to content

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

Closed
Closed
[Tracker] fix training resuming problem when using FP16 in the examples#6552
Labels
Good second issuecontributions-welcome

Description

@sayakpaul
sayakpaul
opened on Jan 12, 2024

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.")

#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:

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.

Activity

  1. sayakpaul
    added
    Good second issue
    contributions-welcome
    on Jan 12, 2024

  2. a-r-r-o-w commented on Jan 12, 2024

    @a-r-r-o-w
    a-r-r-o-w
    on Jan 12, 2024
    Contributor

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

  3. steverhoades commented on Jan 12, 2024

    @steverhoades
    steverhoades
    on Jan 12, 2024
    Contributor

    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.

  4. steverhoades commented on Jan 13, 2024

    @steverhoades
    steverhoades
    on Jan 13, 2024
    Contributor

    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]).
  ...

  5. sayakpaul commented on Jan 13, 2024

    @sayakpaul
    sayakpaul
    on Jan 13, 2024
    MemberAuthor

    @linoytsaban ^

  6. steverhoades commented on Jan 13, 2024

    @steverhoades
    steverhoades
    on Jan 13, 2024
    Contributor

    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"
  7. sayakpaul
    changed the title [-][Examples] fix training resuming problem when using FP16[/-] [+][Tracker] fix training resuming problem when using FP16 in the examples[/+] on Jan 29, 2024

  8. pedrogengo commented on Jan 30, 2024

    @pedrogengo
    pedrogengo
    on Jan 30, 2024
    Contributor

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

  9. sriagastya commented on Feb 4, 2024

    @sriagastya
    sriagastya
    on Feb 4, 2024
    Contributor

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

  10. sayakpaul commented on Feb 4, 2024

    @sayakpaul
    sayakpaul
    on Feb 4, 2024
    MemberAuthor

    Please feel free to.

  11. ParagEkbote commented on May 9, 2025

    @ParagEkbote
    ParagEkbote
    on May 9, 2025
    Contributor

    Has this issue been resolved or are there any pending scripts I could contribute to fix?

    cc: @sayakpaul

  12. sayakpaul commented on May 10, 2025

    @sayakpaul
    sayakpaul
    on May 10, 2025
    MemberAuthor

    You can take an example where the fix hasn't been propagated yet and open a PR.

  13. ParagEkbote commented on May 10, 2025

    @ParagEkbote
    ParagEkbote
    on May 10, 2025
    Contributor

    I believe that the fix has been propagated to Advanced LoRA trainer SD v1.5 in PR 7018 which is the final script in the tracker. Are there more scripts in the examples folder that need the fix to be applied?

    cc: @sayakpaul

  14. sayakpaul commented on May 10, 2025

    @sayakpaul
    sayakpaul
    on May 10, 2025
    MemberAuthor

    How about https://github.com/huggingface/diffusers/blob/main/examples/advanced_diffusion_training/train_dreambooth_lora_sd15_advanced.py?

  15. ParagEkbote commented on May 10, 2025

    @ParagEkbote
    ParagEkbote
    on May 10, 2025
    Contributor

    So, PR 7018 did not propagate the fix and the script is still open for contribution, right?

  16. sayakpaul commented on May 10, 2025

    @sayakpaul
    sayakpaul
    on May 10, 2025
    MemberAuthor

    My bad, sorry. It did propagate it. Closing this issue.

  17. sayakpaul
    closed this as completedon May 10, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    Good second issuecontributions-welcome

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions