SDXL dreambooth can't be resumed from a checkpoint at fp16 training

[epi-morphism]

Describe the bug

train_dreambooth_lora_sdxl.py can't be resumed from a checkpoint using fp16. The log error is Attempting to unscale FP16 gradients.

This is a big blocker from being able to train on the free colab tier since you need fp16 to fit in vram, but also need to resume from checkpoints since it can hit a timeout at any moment.

Reproduction

Reproduce with: https://colab.research.google.com/drive/15woNcXcpsa3GDGk6cmDtIL2V8zRtOOj3

Logs

No response

System Info

latest diffusers, system is whatever is on colab (see linked colab above)

Who can help?

@patrickvonplaten @sayakpaul

[sayakpaul]

Cc: @muellerzr @SunMarc

[danieltanhx]

temporary solution is to brute force into /usr/local/lib/python3.10/dist-packages/torch/cuda/amp/grad_scaler.py change line 284 from "optimizer_state["found_inf_per_device"] = self.unscale_grads(optimizer, inv_scale, found_inf, False)" to "optimizer_state["found_inf_per_device"] = self.unscale_grads(optimizer, inv_scale, found_inf, True)"

[SunMarc]

Hi @danieltanhx , indeed these is an issue with the optimizer where the gradient are somehow calculated in dtype torch.float16 instead of torch.float32 after resuming the training. We will work on fixing this asap cc @muellerzr .

[patrickvonplaten]

Let's wait for the next accelerate release here then

[epi-morphism]

temporary solution is to brute force into /usr/local/lib/python3.10/dist-packages/torch/cuda/amp/grad_scaler.py change line 284 from "optimizer_state["found_inf_per_device"] = self.unscale_grads(optimizer, inv_scale, found_inf, False)" to "optimizer_state["found_inf_per_device"] = self.unscale_grads(optimizer, inv_scale, found_inf, True)"

Have confirmed that this silences the error and lets it continue from a checkpoint. However the end trained lora result doesn't look that good so I suspect something is still off, i.e. not quite a working solution. As patrick said I guess we'll wait until the next accelerate release.

[SunMarc]

After investigation, it seems like it is an issue on diffusers side.
When we resume the checkpoint, we load back the unet lora weights. It uses successively the following functions load_model_hook, load_lora_into_unet and load_attn_procs. In load_attn_procs, the entire unet with lora weight will be converted to the dtype of the unet. The problem is that in the script, we cast all non-trainable weigths to fp16 or bf16 for mixed precision training setup (vae, non-lora text_encoder and non-lora unet).

# For mixed precision training we cast all non-trainable weigths (vae, non-lora text_encoder and non-lora unet) to half-precision
# as these weights are only used for inference, keeping weights in full precision is not required.
weight_dtype = torch.float32
if accelerator.mixed_precision == "fp16":
    weight_dtype = torch.float16
elif accelerator.mixed_precision == "bf16":
    weight_dtype = torch.bfloat16

# Move unet, vae and text_encoder to device and cast to weight_dtype
unet.to(accelerator.device, dtype=weight_dtype)

So in the end, the lora layers end up being converted also in load_attn_procs. See code below or the link to the code source.

# set lora layers
for target_module, lora_layer in lora_layers_list:
    target_module.set_lora_layer(lora_layer)
self.to(dtype=self.dtype, device=self.device)

To solve this issue, we can either modify load_attn_procs to set the lora_layers after this line
self.to(dtype=self.dtype, device=self.device) or convert the lora_layers back to fp32 in the script.
LMK what which option you prefer @patrickvonplaten , it will do a PR to fix that.

[patrickvonplaten]

It would be nice if we could convert the lora_layers back to fp32 to solve the issue. IMO the difference in precision between layers is a idiosyncrasy of the example scripts and should not influence how the load_attn_procs work for inference.

It would be super nice to open a PR fixing it only in the examples scripts

[yuxu915]

@SunMarc @patrickvonplaten hi, really thanks for your investigation and solution.
But I still need more specific details on how to make the changes.
Could you please provide step-by-step instructions or a code snippet to help me understand how to apply your solution?
Your assistance would be greatly appreciated.

[yuxu915]

@SunMarc @patrickvonplaten hi, really thanks for your investigation and solution.
I modify the code as: unet.to(accelerator.device, dtype=torch.float32) to fix the bug, however, when I check the validation results, it seems that training didn't start from the checkpoint I specified; instead, it started from the beginning.
could you help solving the problem? Thank you so much.

[yuxu915]

temporary solution is to brute force into /usr/local/lib/python3.10/dist-packages/torch/cuda/amp/grad_scaler.py change line 284 from "optimizer_state["found_inf_per_device"] = self.unscale_grads(optimizer, inv_scale, found_inf, False)" to "optimizer_state["found_inf_per_device"] = self.unscale_grads(optimizer, inv_scale, found_inf, True)"

Have confirmed that this silences the error and lets it continue from a checkpoint. However the end trained lora result doesn't look that good so I suspect something is still off, i.e. not quite a working solution. As patrick said I guess we'll wait until the next accelerate release.

hi, @epi-morphism , have you finally solved the problem? Would you like to share your solution?

[epi-morphism]

temporary solution is to brute force into /usr/local/lib/python3.10/dist-packages/torch/cuda/amp/grad_scaler.py change line 284 from "optimizer_state["found_inf_per_device"] = self.unscale_grads(optimizer, inv_scale, found_inf, False)" to "optimizer_state["found_inf_per_device"] = self.unscale_grads(optimizer, inv_scale, found_inf, True)"

Have confirmed that this silences the error and lets it continue from a checkpoint. However the end trained lora result doesn't look that good so I suspect something is still off, i.e. not quite a working solution. As patrick said I guess we'll wait until the next accelerate release.

hi, @epi-morphism , have you finally solved the problem? Would you like to share your solution?

Nope, just waiting

[yuxu915]

I modify the code as: unet.to(accelerator.device, dtype=torch.float32)

As mentioned above, I 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 , the error is fixed, however, training seems start at very beginning rather than the designated checkpoint.

[SunMarc]

I was able to fix the prior version of the file, but the new commit of the example with the new lora layers returns another issue when using mixed_precision + resuming checkpoint. I will investigate more later.

  File "/home/marc/anaconda3/envs/accelerate/lib/python3.11/site-packages/torch/cuda/amp/grad_scaler.py", line 372, 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.