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

@patrickvonplaten

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

@muellerzr

Cc: @muellerzr @SunMarc

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

@danieltanhx

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 .

Let's wait for the next accelerate release here then

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.

@patrickvonplaten

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.

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

@SunMarc

@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.

@SunMarc

@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.

@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?

@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

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.

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.

@sayakpaul

cc @sayakpaul can you take a look here?

Yup works with PT 2.0 as well: https://wandb.ai/xin-sayak/dreambooth-lora-sd-xl/runs/ketnbcdw.

Hi, what are the colab notebooks intended to demonstrate? I just see that it looks like you get a training run going (the wandb logs are private in case you intended for there to be information there)

@sayakpaul

hi, @sayakpaul @epi-morphism , 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.

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.

still an issue, not stale

@SunMarc

Have you tried this with the updated version of the script? Cc: @SunMarc

@SunMarc

Have you tried this with the updated version of the script? Cc: @SunMarc

Tried the latest version just now. Can't say if or if it hasn't been fixed since it hits this error

ValueError: Query/Key/Value should either all have the same dtype, or (in the quantized case) Key/Value should have dtype torch.int32
  query.dtype: torch.float32
  key.dtype  : torch.float16
  value.dtype: torch.float16

Reproduce with colab at the top.

@epi-morphism

Hi @epi-morphism, this looks like an issue on xformers. See related thread in accelerate. Try to download and compile xformers from source.

@SunMarc

Have you tried this with the updated version of the script? Cc: @SunMarc @sayakpaul

Ok I fixed the above (by not using memory efficient attention) and just tried it again. The issue is the same as in the original post i.e.

ValueError: Attempting to unscale FP16 gradients.

when trying to resume with fp16. So the issue is not stale, still a bug.
See exact logs and reproduce with this https://colab.research.google.com/drive/1hW5vWnhXjPyHAuq9p9LdhUcej7GFt6-u

@SunMarc

So, it's the training resume step where you're facing problems, yeah? Again, cc: @SunMarc.

This needs to be solved first: #6510.

Tracked this down to the norm layers.

Tracing the types they go from float16 -> float32

diffusers/src/diffusers/models/attention.py

Lines 299 to 317 in dd63168

    
           if self.use_ada_layer_norm: 
        
               norm_hidden_states = self.norm1(hidden_states, timestep) 
        
           elif self.use_ada_layer_norm_zero: 
        
               norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1( 
        
                   hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype 
        
               ) 
        
           elif self.use_layer_norm: 
        
               norm_hidden_states = self.norm1(hidden_states) 
        
           elif self.use_ada_layer_norm_continuous: 
        
               norm_hidden_states = self.norm1(hidden_states, added_cond_kwargs["pooled_text_emb"]) 
        
           elif self.use_ada_layer_norm_single: 
        
               shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( 
        
                   self.scale_shift_table[None] + timestep.reshape(batch_size, 6, -1) 
        
               ).chunk(6, dim=1) 
        
               norm_hidden_states = self.norm1(hidden_states) 
        
               norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa 
        
               norm_hidden_states = norm_hidden_states.squeeze(1) 
        
           else: 
        
               raise ValueError("Incorrect norm used")

https://github.com/huggingface/diffusers/blob/dd63168319d4cd3c5f5951daaf2a37a26d2f9b73/src/diffusers/models/attention.py#L350C13-L361

Have been solved with: #6514. So, I am actually going to close this.

Navigation Menu

Repository navigation

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

Describe the bug

Reproduction

Logs

System Info

Who can help?

12 remaining items

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

	if self.use_ada_layer_norm:
	norm_hidden_states = self.norm1(hidden_states, timestep)
	elif self.use_ada_layer_norm_zero:
	norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(
	hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype
	)
	elif self.use_layer_norm:
	norm_hidden_states = self.norm1(hidden_states)
	elif self.use_ada_layer_norm_continuous:
	norm_hidden_states = self.norm1(hidden_states, added_cond_kwargs["pooled_text_emb"])
	elif self.use_ada_layer_norm_single:
	shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
	self.scale_shift_table[None] + timestep.reshape(batch_size, 6, -1)
	).chunk(6, dim=1)
	norm_hidden_states = self.norm1(hidden_states)
	norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa
	norm_hidden_states = norm_hidden_states.squeeze(1)
	else:
	raise ValueError("Incorrect norm used")

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

Description

Describe the bug

Reproduction

Logs

System Info

Who can help?

Activity

sayakpaul commented on Sep 13, 2023

danieltanhx commented on Sep 20, 2023

SunMarc commented on Sep 20, 2023

patrickvonplaten commented on Sep 20, 2023

epi-morphism commented on Sep 24, 2023

SunMarc commented on Sep 29, 2023

patrickvonplaten commented on Oct 2, 2023

yuxu915 commented on Oct 15, 2023

yuxu915 commented on Oct 16, 2023

yuxu915 commented on Oct 17, 2023

epi-morphism commented on Oct 17, 2023

yuxu915 commented on Oct 17, 2023

SunMarc commented on Oct 18, 2023

patrickvonplaten commented on Oct 23, 2023

12 remaining items

sayakpaul commented on Oct 26, 2023

LWprogramming commented on Oct 26, 2023

yuxu915 commented on Nov 16, 2023

github-actions commented on Dec 26, 2023

epi-morphism commented on Jan 3, 2024

sayakpaul commented on Jan 4, 2024

epi-morphism commented on Jan 8, 2024

SunMarc commented on Jan 9, 2024

epi-morphism commented on Jan 9, 2024

sayakpaul commented on Jan 10, 2024

sayakpaul commented on Jan 10, 2024

nbardy commented on Jan 17, 2024

sayakpaul commented on Jan 17, 2024

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions