Support dynamic LoRA loading with `torch.compile` model

[apolinario]

Is your feature request related to a problem? Please describe.
Would be great to be able to load a LoRA to a model compiled with torch.compile

Describe the solution you'd like.
Do load_lora_weights with a compiled pipe (ideally without triggering recompilation)

Currently, running this code:

import torch
from diffusers import DiffusionPipeline

device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
pipe = pipe.to(device)
pipe.transformer = pipe.transformer.to(memory_format=torch.channels_last)
pipe.transformer = torch.compile(pipe.transformer, mode="reduce-overhead")

pipe.load_lora_weights("multimodalart/flux-tarot-v1")

It errors:

Loading adapter weights from state_dict led to unexpected keys not found in the model:  ['single_transformer_blocks.0.attn.to_k.lora_A.default_3.weight', 'single_transformer_blocks.0.attn.to_k.lora_B.default_3.weight', 'single_transformer_blocks.0.attn.to_q.lora_A.default_3.weight', 'single_transformer_blocks.0.attn.to_q.lora_B.default_3.weight',

When compiled, the state dict of a model seem to add a _orig_mod prefix to all keys

odict_keys(['_orig_mod.time_text_embed.timestep_embedder.linear_1.weight',...

Describe alternatives you've considered.
An alternative is to fuse the LoRA into the model and then compile, however this does not allow for hot swapping LoRAs (as a new pipeline and a new compilation is needed for every LoRA)

Additional context.
This seems to have been achieved by @chengzeyi , author of the now paused https://github.com/chengzeyi/stable-fast , however it seems to be part of the non-open source FAL optimized inference (however if you'd like to contribute this upstream, feel free!)

[apolinario]

Some references:

• PEFT seems to support torch_compile for some operations https://huggingface.co/docs/peft/main/en/developer_guides/torch_compile
• When compiled, the state dict of a model seem to add a _orig_mod prefix to all keys

odict_keys(['_orig_mod.time_text_embed.timestep_embedder.linear_1.weight', '_orig_mod.time_text_embed.timestep_embedder.linear_1.bias'...

[sayakpaul]

@apolinario thanks for the detailed thread. Would be also nice to include a small snippet that we could quickly verify. And yes, ccing @BenjaminBossan for his comments here.

I assume we could just add checks in load_lora_weights() just before when we pass the adapter weights to the underlying model:

diffusers/src/diffusers/loaders/lora_pipeline.py

Line 1218 in 4cfb216

incompatible_keys = set_peft_model_state_dict(transformer, state_dict, adapter_name)

[BenjaminBossan]

I have not dealt with hot-swapping LoRA weights on compiled models in PEFT, so I'm not surprised that it doesn't work out of the box. The PEFT experiments with compiled models all apply the compilation step after loading the PEFT weights. Maybe it's as easy as remapping the state dict and then calling set_peft_model_state_dict, but I wouldn't be surprised if there are more pitfalls related to torch.compile implementation details.

I'll make a note to look into this on the PEFT side when I have a bit of extra time. But do let me know if you make any progress.

[apolinario]

Would be also nice to include a small snippet that we could quickly verify

Added!

[sayakpaul]

Applying this change seems to work:
0e7204a

However, I get the following when doing torch.compile() with mode="max-autotune", fullgraph=True:

skipping cudagraphs due to skipping cudagraphs due to cpu device (arg25_1). Found from : 
   File "/fsx/sayak/diffusers/src/diffusers/models/transformers/transformer_flux.py", line 494, in forward
    encoder_hidden_states, hidden_states = block(
  File "/fsx/sayak/diffusers/src/diffusers/models/transformers/transformer_flux.py", line 165, in forward
    norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb)
  File "/fsx/sayak/diffusers/src/diffusers/models/normalization.py", line 137, in forward
    emb = self.linear(self.silu(emb))
  File "/fsx/sayak/miniconda3/envs/diffusers/lib/python3.9/site-packages/peft/tuners/lora/layer.py", line 509, in forward
    result = result + lora_B(lora_A(dropout(x))) * scaling

[apolinario]

Cool, 0e7204a works here too for reduce-overhead

However it triggers a recompilation when you load/swap a LoRA:

import torch
from diffusers import DiffusionPipeline

device = "cuda" if torch.cuda.is_available() else "cpu"

pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
pipe = pipe.to(device)
pipe.transformer = pipe.transformer.to(memory_format=torch.channels_last)
pipe.transformer = torch.compile(pipe.transformer, mode="reduce-overhead")

prompt = "a photo of an astronaut riding a horse on mars"
#This will compile for the first time
image = pipe(prompt).images[0]

pipe.load_lora_weights("multimodalart/flux-tarot-v1")
prompt = "a photo of an astronaut riding a horse on mars, tarot card"
#This will re-compile
image = pipe(prompt).images[0]

pipe.unload_lora_weights()
pipe.load_lora_weights("davisbro/half_illustration")
prompt = "In the style of TOK, a photo of an astronaut riding a horse on mars"
#This will re-compile again
image = pipe(prompt).images[0]

[sayakpaul]

Oh okay. Let's perhaps discuss this with the PyTorch team.

[anijain2305]

Recompiles are coming from this

V0827 10:44:01.036000 3046403 torch/_dynamo/guards.py:2796] [0/1] [__recompiles]     triggered by the following guard failure(s):
V0827 10:44:01.036000 3046403 torch/_dynamo/guards.py:2796] [0/1] [__recompiles]     - 0/0: ___check_type_id(L['self']._modules['transformer_blocks']._modules['0']._modules['norm1']._modules['linear'], 81088688)  # emb = self.linear(self.silu(emb))  # diffusers/src/diffusers/models/normalization.py:137 in forward

On the first compile, TorchDynamo guarded on the id(type(torch.nn.Linear)). But loading weights has somehow changed the type of the linear layer (or atleast the id of the type of class). Do you know if loading weights of the linear changes its class type (of if we are dynamically creating a new class)?

[sayakpaul]

Yeah it does.

When we're hitting

diffusers/src/diffusers/loaders/lora_pipeline.py

Line 1217 in 4cfb216

inject_adapter_in_model(lora_config, transformer, adapter_name=adapter_name)

The codepath in peft is

https://github.com/huggingface/peft/blob/850eeb5c3a5cf692f5612c7c733b13fde184e05d/src/peft/mapping.py#L223

To get around the problem, we can likely just fuse the adapter modules into the respective base model modules to prevent recompilation but this is somewhat a bit restrictive in terms of a smoother UX. Perhaps @apolinario can explain this a bit better. So, we wanted to see if there's any alternative we could try here.

Cc: @BenjaminBossan

[anijain2305]

No easy alternative I can think of as of now. torch.compile is adding a guard on the type. In this case, it seems that the type change is benign, but in general the type change requires a recompilation.

Is it possible to somehow ensure that type of the class is same as before after all the weights have been loaded? If we can do that, there wont be any recompiles. A bad way would be to monkeypatch the __class__ itself to the old type.

[Chillee]

If you're loading a Lora, it's very reasonable to have a recompilation, no? The actual operations are different. The question here is whether it needs a recompilation upon swapping a new LoRA.

The other thing is that if you call model.compile() (as opposed to torch.compile(model)), the state dict won't be modified.

[apolinario]

If you're loading a Lora, it's very reasonable to have a recompilation, no? The actual operations are different. The question here is whether it needs a recompilation upon swapping a new LoRA.

Agreed. I think loading the first LoRA triggering recompilation is probably fine. However, hot swapping different LoRAs and having it not recompile would be very good to allow for dynamic apps that change LoRAs to benefit from compilation performances (as waiting for it to compile every time LoRAs are swapped would not allow for a live swapping application)

[anijain2305]

I see. In that case, can someone run TORCH_LOGS="guards,recompiles" and share the log? For some reason, my run fails at

pipe.unload_lora_weights()
pipe.load_lora_weights("davisbro/half_illustration")

[sayakpaul]

I am getting you the log but recompilation upon swapping or loading a new LoRA also seems reasonable to me:

• The new LoRA may have a different rank
• It may have different target layers than the one initially loaded

Both of these combined should lead to different operations, I'd imagine.