Fix: training resume from fp16 for SDXL Consistency Distillation

[asrimanth]

What does this PR do?

Part of #6552 for SDXL Consistency Distillation

Before submitting

• Did you read the contributor guideline?
• Did you read our philosophy doc (important for complex PRs)?
• Was this discussed/approved via a GitHub issue? Yes, here's the discussion
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

What to test?

First run this:

accelerate launch train_lcm_distill_lora_sdxl.py \
  --pretrained_teacher_model="stabilityai/stable-diffusion-xl-base-1.0"  \
  --pretrained_vae_model_name_or_path="madebyollin/sdxl-vae-fp16-fix" \
  --output_dir="pokemons-lora-lcm-sdxl" \
  --mixed_precision="fp16" \
  --dataset_name="lambdalabs/pokemon-blip-captions" \
  --resolution=512 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=1 --gradient_checkpointing \
  --use_8bit_adam \
  --lora_rank=16 \
  --learning_rate=1e-4 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=4 --checkpointing_steps=2 --checkpoints_total_limit=2 \
  --seed="0"

To resume training, run the following command:

accelerate launch train_lcm_distill_lora_sdxl.py \
  --pretrained_teacher_model="stabilityai/stable-diffusion-xl-base-1.0"  \
  --pretrained_vae_model_name_or_path="madebyollin/sdxl-vae-fp16-fix" \
  --output_dir="pokemons-lora-lcm-sdxl" \
  --mixed_precision="fp16" \
  --dataset_name="lambdalabs/pokemon-blip-captions" \
  --resolution=512 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=1 --gradient_checkpointing \
  --use_8bit_adam \
  --lora_rank=16 \
  --learning_rate=1e-4 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=8 --checkpointing_steps=2 --checkpoints_total_limit=2 \
  --seed="0" \
  --resume_from_checkpoint="latest"

Who can review?

@sayakpaul

[sayakpaul]

Thanks for your contribution. Could we please reduce the number of steps to quickly check this?

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[asrimanth]

Do you mean reducing the max_train_steps or the checkpointing_steps? FYI, the above-mentioned tests work for me in my local machine. Also, how do I fix the code quality check?

[sayakpaul]

Do you mean reducing the max_train_steps or the checkpointing_steps? FYI, the above-mentioned tests work for me in my local machine. Also, how do I fix the code quality check?

Yeah. I am sure the above command will work but will take longer to validate. The effectivity can essentially be tested with far fewer steps, as I did here. Also, please run the code quality linter by running make style && make quality.

[asrimanth]

I've pushed a commit for the code style, and I've edited the test command to save a checkpoint for fewer number of steps, as per your example.

[sayakpaul]

I don't think I made myself super clear. Sorry.

Your example commands still have max_train_steps set to 3000. Why though? Here, I have it set to 4 8 respectively along with checkpointing_steps and checkpoints_total_limit set accordingly. I don't have any unnecessary arguments here such as report_to.

Can we please keep the example commands super lean?

I hope I made myself more clear this time.

[asrimanth]

Yeah, makes sense! I've made the changes so that the training command is leaner and more efficient.

[sayakpaul]

I just tried it and I am facing shape mismatch problems while running the second command.

[asrimanth]

Updates: I just got a Runpod machine with higher VRAM for testing. I was able to test the updated script with lower settings. I also made some changes to the code and it works on my local machine. Please test it out and let me know.

[sayakpaul]

Thank you! But now, I am seeing some state dict key mismtaches when I am resuming training. Could you look into that?

[asrimanth]

Pushed a new commit to fix the missing keys issue. Please have a look and let me know.

[sayakpaul]

We don't need to change this.

[asrimanth]

Forgot to remove it after local testing. Reverted it to None and pushed a commit. Sorry for the inconvenience.

[sayakpaul]

Thanks so much for iterating. I have checked and it works: https://colab.research.google.com/gist/sayakpaul/fd2e863c9911031ad01fa9cf6863a5da/scratchpad.ipynb.

Will merge the PR once the test suite passes.

[asrimanth]

Thank you for the consistent feedback. Happy to contribute to HuggingFace.