resume_from_checkpoint training fail in train_dreambooth_lora_sdxl.py

[yuxu915]

Describe the bug

I trained dreambooth with lora and sd-xl for 1000 steps, then I try to continue traning resume from the 500th step, however, it seems like the training starts without the 1000's checkpoint, i.e. it starts from the beginning. Training scripts are as below.
btw. I fix another resume bug as advised in #5004.

could you please help it?

Reproduction

export MODEL_NAME="/data/model/stable-diffusion-xl-base-1.0/stable-diffusion-xl-base-1.0"
export INSTANCE_DIR="/data/datasets/image_instance/chair_crop"
export OUTPUT_DIR="/data/lora-xl/diffusers/examples/dreambooth/lora-trained-xl_1e-6_step1500_chair_crop_per25"
export VAE_PATH="/data/model/stable-diffusion-xl-base-1.0/sdxl-vae-fp16-fix"

accelerate launch train_dreambooth_lora_sdxl.py \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --instance_data_dir=$INSTANCE_DIR \
  --pretrained_vae_model_name_or_path=$VAE_PATH \
  --output_dir=$OUTPUT_DIR \
  --mixed_precision="fp16" \
  --instance_prompt="a photo of sks chair" \
  --resolution=1024 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --learning_rate=1e-5 \
  --report_to="wandb" \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=2700 \
  --validation_prompt="A photo of sks chair on grass" \
  --validation_epochs=25 \
  --seed="0" \
  --checkpointing_steps=500 \
  --resume_from_checkpoint="/data/xuyu/lora-xl/diffusers/examples/dreambooth/lora-trained-xl_1e-6_step1500_chair_crop_per25/checkpoint-1000"

Logs

No response

System Info

• diffusers version: 0.22.0.dev0
• Platform: Linux-3.10.0-1160.95.1.el7.x86_64-x86_64-with-glibc2.17
• Python version: 3.10.13
• PyTorch version (GPU?): 2.1.0+cu121 (True)
• Huggingface_hub version: 0.17.3
• Transformers version: 4.34.0
• Accelerate version: 0.23.0
• xFormers version: not installed
• Using GPU in script?: yes
• Using distributed or parallel set-up in script?: parallel

Who can help?

@patrickvonplaten @sayakpaul @yiyixuxu @dn

[sayakpaul]

Can you provide a reproducible Colab Notebook?

[yuxu915]

hi @sayakpaul ,
I appreciate your request for a Colab notebook. However, I've been running the code locally and haven't made any substantial changes to the default configuration provided in the repository. As a result, I don't have a specific Colab notebook available for this code.
If you need assistance with any specific sections of the code or have questions about its functionality, I'd be happy to provide guidance based on the local setup.
Really thank you for your understanding.

[sayakpaul]

Unfortunately, without something for us to verify concretely, we cannot proceed further. As you can see here, there's already a test that checks for whether or training run for DreamBooth LoRA can be resumed successfully:

diffusers/examples/test_examples.py

Line 1401 in b4cbbd5

def test_dreambooth_lora_checkpointing_checkpoints_total_limit_removes_multiple_checkpoints(self):

[yuxu915]

Can you provide a reproducible Colab Notebook?

hi, @sayakpaul , thanks for responce, reproducible Colab Notebook is as follows:
https://colab.research.google.com/drive/15woNcXcpsa3GDGk6cmDtIL2V8zRtOOj3
please notice that in order to fix the error Attempting to unscale FP16 gradients., we modify the code from unet.to(accelerator.device, dtype=weight_dtype) to unet.to(accelerator.device, dtype=torch.float32) in train_dreambooth_lora_sdxl.py in line 719.

[yuxu915]

hi, @sayakpaul ,another bug is that when resume training from a middle lora checkpoint, it stops update the model( i.e. checkpionts remain the same as the start one), have you ever meet this bug? really thanks.

[sayakpaul]

Please open a new issue for this. I don't think it is really related to the current issue of yours.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.