train_dreambooth_lora.py resume from checkpoint errors.

[JaLnYn]

Describe the bug

I get the following issue when trying to resume from checkpoint.
KeyError: 'unet.down_blocks.0.attentions.0.transformer_blocks.0.attn1.processor'
There was also a naming issue where I had to change pytorch_lora_weights.bin to pytorch_model.bin in the checkpoint folders

Reproduction

accelerate launch train_dreambooth_lora.py
--pretrained_model_name_or_path=$MODEL_NAME
--instance_data_dir=$INSTANCE_DIR
--output_dir=$OUTPUT_DIR
--instance_prompt="a photo of sks dog"
--resolution=512
--train_batch_size=1
--gradient_accumulation_steps=1
--checkpointing_steps=100
--learning_rate=1e-4
--report_to="wandb"
--lr_scheduler="constant"
--lr_warmup_steps=0
--max_train_steps=500
--validation_prompt="A photo of sks dog in a bucket"
--validation_epochs=50
--seed="0"
--push_to_hub
--resume_from_checkpoint="" \

Logs

(most recent call last):
  File "/home/rip/git/diffusers/examples/dreambooth/train_dreambooth_lora.py", line 1093, in <module>
    main(args)
  File "/home/rip/git/diffusers/examples/dreambooth/train_dreambooth_lora.py", line 881, in main
    accelerator.load_state(os.path.join(args.output_dir, path))
  File "/home/rip/.local/lib/python3.11/site-packages/accelerate/accelerator.py", line 2396, in load_state
    load_accelerator_state(
  File "/home/rip/.local/lib/python3.11/site-packages/accelerate/checkpointing.py", line 140, in load_accelerator_state
    models[i].load_state_dict(torch_mod, **load_model_func_kwargs)
  File "/home/rip/.local/lib/python3.11/site-packages/torch/nn/modules/module.py", line 2027, in load_state_dict
    load(self, state_dict)
  File "/home/rip/.local/lib/python3.11/site-packages/torch/nn/modules/module.py", line 2009, in load
    module._load_from_state_dict(
  File "/home/rip/.local/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1909, in _load_from_state_dict
    hook(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
  File "/home/rip/.local/lib/python3.11/site-packages/torch/nn/modules/module.py", line 69, in __call__
    return self.hook(module, *args, **kwargs)
  File "/home/rip/git/diffusers/src/diffusers/loaders.py", line 84, in map_from
    new_key = key.replace(replace_key, f"layers.{module.rev_mapping[replace_key]}")
                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
KeyError: 'unet.down_blocks.0.attentions.0.transformer_blocks.0.attn1.processor'

System Info

• diffusers version: 0.17.0.dev0
• Platform: Linux-5.10.16.3-microsoft-standard-WSL2-x86_64-with-glibc2.35
• Python version: 3.11.3
• PyTorch version (GPU?): 2.0.0+cu117 (True)
• Huggingface_hub version: 0.14.0
• Transformers version: 4.28.1
• Accelerate version: 0.18.0
• xFormers version: not installed
• Using GPU in script?: yes
• Using distributed or parallel set-up in script?: no

[JaLnYn]

accelerator.load_state(os.path.join(args.output_dir, path)) is not the correct approach for loading lora wieghts. Attempting to modify unet variable to load weights
added

# accelerator.load_state(os.path.join(args.output_dir, path))
unet.load_attn_procs(os.path.join(args.output_dir, path))

to examples/dreambooth/train_dreambooth_lora.py:881 and get this error:

Traceback (most recent call last):
  File "/home/rip/git/diffusers/examples/dreambooth/train_dreambooth_lora.py", line 1097, in <module>
    main(args)
  File "/home/rip/git/diffusers/examples/dreambooth/train_dreambooth_lora.py", line 961, in main
    optimizer.step()
  File "/home/rip/.local/lib/python3.11/site-packages/accelerate/optimizer.py", line 134, in step
    self.scaler.step(self.optimizer, closure)
  File "/home/rip/.local/lib/python3.11/site-packages/torch/cuda/amp/grad_scaler.py", line 368, in step
    assert len(optimizer_state["found_inf_per_device"]) > 0, "No inf checks were recorded for this optimizer."
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError: No inf checks were recorded for this optimizer.

I have no idea what I'm doing someone please help

[fractaldna22]

im getting the same error after it worked the first time, trained with lora layers as the processor, saved the pipeline as normal, then the next time i try to load it says key error.

  text_encoder = CLIPTextModel.from_pretrained(
        args.pretrained_model_name_or_path,
        subfolder="text_encoder",
        revision=args.revision,
    )
    vae = AutoencoderKL.from_pretrained(
        args.pretrained_model_name_or_path,
        subfolder="vae",
        revision=args.revision,
    )
    unet = UNet2DConditionModel.from_pretrained(
        args.pretrained_model_name_or_path,
        subfolder="unet",
        revision=args.revision,
        torch_dtype=torch.float32
    )
######################################
File "/content/diffusers/src/diffusers/models/modeling_utils.py", line 602, in from_pretrained
    if empty_state_dict[param_name].shape != param.shape:
KeyError: 'down_blocks.0.attentions.0.transformer_blocks.0.attn1.processor.to_q_lora.down.weight'

I take this to mean that the last 2 hours of training i did is now trashed.

[patrickvonplaten]

I cannot reproduce this error cc @sayakpaul maybe also for visibility

[JaLnYn]

I think the issue is related to me setting the fp16 in the accelerate config. I'm also not sure how to cc...

[JaLnYn]

I cannot reproduce this error cc @sayakpaul maybe also for visibility

This line should require, optimizer and other accelerate.load_state files, Using the default lora saves, it only saves the lora_wieghts.bin. I think if you're using the newest train_dreambooth_lora file you should run into this error
876 accelerator.load_state(os.path.join(args.output_dir, path))

[sayakpaul]

Hmm, I think the resuming logic needs to be modified for LoRA training examples. I don't think it's related to FP16. Cc: @williamberman any suggestions on how to best approach this?

[JaLnYn]

Hmm, I think the resuming logic needs to be modified for LoRA training examples. I don't think it's related to FP16. Cc: @williamberman any suggestions on how to best approach this?

Sorry for the fp16 comment I forgot to add context

I've removed the accelerate.loade_state line because I think it's wrong.

I attempted to add
unet.load_attn_procs(path to weights):

 709     if args.resume_from_checkpoint:
 710         if args.resume_from_checkpoint != "latest":
 711             path = os.path.basename(args.resume_from_checkpoint)
 712         else:
 713             # Get the mos recent checkpoint
 714             dirs = os.listdir(args.output_dir)
 715             dirs = [d for d in dirs if d.startswith("checkpoint")]
 716             dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
 717             path = dirs[-1] if len(dirs) > 0 else None
 718
 719         if path is None:
 720             accelerator.print(
 721                 f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run."
 722             )
 723             args.resume_from_checkpoint = None
 724         else:
 725             unet.load_attn_procs(os.path.join(args.output_dir, path))
 726     else:   
 727         unet_lora_attn_procs = {}
 728         for name in unet.attn_processors.keys():
 729             cross_attention_dim = None if name.endswith("attn1.processor") else unet.config.cross_attention_dim
 730             if name.startswith("mid_block"):
 731                 hidden_size = unet.config.block_out_channels[-1]
 732             elif name.startswith("up_blocks"):
 733                 block_id = int(name[len("up_blocks.")])
 734                 hidden_size = list(reversed(unet.config.block_out_channels))[block_id]
 735             elif name.startswith("down_blocks"):
 736                 block_id = int(name[len("down_blocks.")])
 737                 hidden_size = unet.config.block_out_channels[block_id]
 738 
 739             unet_lora_attn_procs[name] = LoRAAttnProcessor(
 740                 hidden_size=hidden_size, cross_attention_dim=cross_attention_dim
 741             )   
 742 
 743         unet.set_attn_processor(unet_lora_attn_procs)
 744     unet_lora_layers = AttnProcsLayers(unet.attn_processors)
 745     accelerator.register_for_checkpointing(unet_lora_layers)
 746 

But get gradient clipping error, which I found might be because of I need to convert fp16 to fp32 after reading a few posts online.

I also attempted to replace with unet.load_attn_procs(path to weights) which gives me a error about inf gradient?

 877     if args.resume_from_checkpoint:
 878         if args.resume_from_checkpoint != "latest":
 879             path = os.path.basename(args.resume_from_checkpoint)
 880         else:
 881             # Get the mos recent checkpoint
 882             dirs = os.listdir(args.output_dir)
 883             dirs = [d for d in dirs if d.startswith("checkpoint")]
 884             dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
 885             path = dirs[-1] if len(dirs) > 0 else None
 886
 887         if path is None:
 888             accelerator.print(
 889                 f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run."
 890             )
 891             args.resume_from_checkpoint = None
 892         else:
 893             accelerator.print(f"Resuming from checkpoint {path}")
 894             # accelerator.load_state(os.path.join(args.output_dir, path))
 895             global_step = int(path.split("-")[1])
 896
 897             unet.load_attn_procs(os.path.join(args.output_dir, path))
 898             # unet = unet.float()
 899             # print(unet.dtype, type(unet))
 900
 901             resume_global_step = global_step * args.gradient_accumulation_steps
 902             first_epoch = global_step // num_update_steps_per_epoch
 903             resume_step = resume_global_step % (num_update_steps_per_epoch * args.gradient_accumulation_steps)
 904

[markrmiller]

I don't think it's related to FP16.

I don't think it is either, and I don't think it's quite as simple as just using the standard way to load lora weights.

I believe when it saves the checkpoints, it writes some kind of semi-bespoke format to store both unet and text_encoder weights in one file I think, and accelerator.load_state has no clue about it. I think there is loader for it on a pipeline, so you can load it for inference, and I've tried to grab the models off that and I've tried to massage the weight names correctly, but some minor detail kept remaining elusive as I went and I so I stopped.

[sayakpaul]

Thanks for sharing your findings. Yeah, we need to look deeper into resuming from a checkpoint when doing LoRA. Let's wait for @williamberman's comments.

[JaLnYn]

It's not a fp16 problem, I converted to fp32 and got dtype mismatch.