Can't torch.compile transformer models that load GGUF via from_single_file

[AstraliteHeart]

Describe the bug

transformer model loaded via GGUF can't be torch.compile(d) and raises torch._dynamo.exc.Unsupported: call_method SetVariable() __setitem__ (UserDefinedObjectVariable(GGUFParameter), ConstantVariable(NoneType: None)) {}

'normal' model loaded from HF for the same pipeline can be torch.compile(d) just fine.

Reproduction

If I load the pipeline from HF model, i.e.

import torch
from diffusers import AuraFlowPipeline

torch.set_float32_matmul_precision("high")
torch._inductor.config.conv_1x1_as_mm = True
torch._inductor.config.coordinate_descent_tuning = True
torch._inductor.config.epilogue_fusion = False
torch._inductor.config.coordinate_descent_check_all_directions = True

pipeline = AuraFlowPipeline.from_pretrained(
    "fal/AuraFlow-v0.3",
    torch_dtype=torch.bfloat16,
).to("cuda")

pipeline.transformer.to(memory_format=torch.channels_last)
pipeline.transformer = torch.compile(pipeline.transformer, mode="max-autotune", fullgraph=True)
pipeline("A cute pony", width=512, height=512, num_inference_steps=5)

I can torch.compile it (and observer better performance).

If I try to load the transformer part from GGUF

import torch
from diffusers import (
    AuraFlowPipeline,
    GGUFQuantizationConfig,
    AuraFlowTransformer2DModel,
)

torch.set_float32_matmul_precision("high")
torch._inductor.config.conv_1x1_as_mm = True
torch._inductor.config.coordinate_descent_tuning = True
torch._inductor.config.epilogue_fusion = False
torch._inductor.config.coordinate_descent_check_all_directions = True

transformer = AuraFlowTransformer2DModel.from_single_file(
    "https://huggingface.co/city96/AuraFlow-v0.3-gguf/blob/main/aura_flow_0.3-Q2_K.gguf",
    quantization_config=GGUFQuantizationConfig(compute_dtype=torch.bfloat16),
    torch_dtype=torch.bfloat16,
)
pipeline = AuraFlowPipeline.from_pretrained(
    "fal/AuraFlow-v0.3",
    torch_dtype=torch.bfloat16,
    transformer=transformer,
).to("cuda")

pipeline.transformer.to(memory_format=torch.channels_last)
pipeline.transformer = torch.compile(pipeline.transformer, mode="max-autotune", fullgraph=True)
pipeline("A cute pony", width=512, height=512, num_inference_steps=5)

it raises an exception (see log).

I am still learning on how torch.compile/dynamo function so its unclear to me if this is just some basic confusion of GGUFParameter wrapping torch.nn.Parameter or if diffusers need to do anything special (or if this is something torch must do better?). I've only tested on AuraFlow but this should be the same for any code using GGUF loading. Happy to continue debugging/raise issue with the torch devs but would appreciate if someone more knowledgable have a look at this.

Logs

Traceback (most recent call last):
  File "test_diff_gguf.py", line 27, in <module>
    pipeline("A cute pony", width=512, height=512, num_inference_steps=5)
  File "/env/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
  File "/env/diffusers/pipelines/aura_flow/pipeline_aura_flow.py", line 555, in __call__
    noise_pred = self.transformer(
  File "/env/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/env/torch/nn/modules/module.py", line 1750, in _call_impl
    return forward_call(*args, **kwargs)
  File "/env/torch/_dynamo/eval_frame.py", line 574, in _fn
    return fn(*args, **kwargs)
  File "/env/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/env/torch/nn/modules/module.py", line 1750, in _call_impl
    return forward_call(*args, **kwargs)
  File "/env/torch/_dynamo/convert_frame.py", line 1380, in __call__
    return self._torchdynamo_orig_callable(
  File "/env/torch/_dynamo/convert_frame.py", line 547, in __call__
    return _compile(
  File "/env/torch/_dynamo/convert_frame.py", line 986, in _compile
    guarded_code = compile_inner(code, one_graph, hooks, transform)
  File "/env/torch/_dynamo/convert_frame.py", line 715, in compile_inner
    return _compile_inner(code, one_graph, hooks, transform)
  File "/env/torch/_utils_internal.py", line 95, in wrapper_function
    return function(*args, **kwargs)
  File "/env/torch/_dynamo/convert_frame.py", line 750, in _compile_inner
    out_code = transform_code_object(code, transform)
  File "/env/torch/_dynamo/bytecode_transformation.py", line 1361, in transform_code_object
    transformations(instructions, code_options)
  File "/env/torch/_dynamo/convert_frame.py", line 231, in _fn
    return fn(*args, **kwargs)
  File "/env/torch/_dynamo/convert_frame.py", line 662, in transform
    tracer.run()
  File "/env/torch/_dynamo/symbolic_convert.py", line 2868, in run
    super().run()
  File "/env/torch/_dynamo/symbolic_convert.py", line 1052, in run
    while self.step():
  File "/env/torch/_dynamo/symbolic_convert.py", line 962, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/env/torch/_dynamo/symbolic_convert.py", line 659, in wrapper
    return inner_fn(self, inst)
  File "/env/torch/_dynamo/symbolic_convert.py", line 1658, in CALL_FUNCTION
    self.call_function(fn, args, {})
  File "/env/torch/_dynamo/symbolic_convert.py", line 897, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
  File "/env/torch/_dynamo/variables/functions.py", line 378, in call_function
    return super().call_function(tx, args, kwargs)
  File "/env/torch/_dynamo/variables/functions.py", line 317, in call_function
    return super().call_function(tx, args, kwargs)
  File "/env/torch/_dynamo/variables/functions.py", line 118, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
  File "/env/torch/_dynamo/symbolic_convert.py", line 903, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "/env/torch/_dynamo/symbolic_convert.py", line 3072, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "/env/torch/_dynamo/symbolic_convert.py", line 3198, in inline_call_
    tracer.run()
  File "/env/torch/_dynamo/symbolic_convert.py", line 1052, in run
    while self.step():
  File "/env/torch/_dynamo/symbolic_convert.py", line 962, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/env/torch/_dynamo/symbolic_convert.py", line 659, in wrapper
    return inner_fn(self, inst)
  File "/env/torch/_dynamo/symbolic_convert.py", line 1748, in CALL_FUNCTION_KW
    self.call_function(fn, args, kwargs)
  File "/env/torch/_dynamo/symbolic_convert.py", line 897, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
  File "/env/torch/_dynamo/variables/functions.py", line 378, in call_function
    return super().call_function(tx, args, kwargs)
  File "/env/torch/_dynamo/variables/functions.py", line 317, in call_function
    return super().call_function(tx, args, kwargs)
  File "/env/torch/_dynamo/variables/functions.py", line 118, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
  File "/env/torch/_dynamo/symbolic_convert.py", line 903, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "/env/torch/_dynamo/symbolic_convert.py", line 3072, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "/env/torch/_dynamo/symbolic_convert.py", line 3198, in inline_call_
    tracer.run()
  File "/env/torch/_dynamo/symbolic_convert.py", line 1052, in run
    while self.step():
  File "/env/torch/_dynamo/symbolic_convert.py", line 962, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/env/torch/_dynamo/symbolic_convert.py", line 659, in wrapper
    return inner_fn(self, inst)
  File "/env/torch/_dynamo/symbolic_convert.py", line 1748, in CALL_FUNCTION_KW
    self.call_function(fn, args, kwargs)
  File "/env/torch/_dynamo/symbolic_convert.py", line 897, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
  File "/env/torch/_dynamo/variables/functions.py", line 378, in call_function
    return super().call_function(tx, args, kwargs)
  File "/env/torch/_dynamo/variables/functions.py", line 317, in call_function
    return super().call_function(tx, args, kwargs)
  File "/env/torch/_dynamo/variables/functions.py", line 118, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
  File "/env/torch/_dynamo/symbolic_convert.py", line 903, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "/env/torch/_dynamo/symbolic_convert.py", line 3072, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "/env/torch/_dynamo/symbolic_convert.py", line 3198, in inline_call_
    tracer.run()
  File "/env/torch/_dynamo/symbolic_convert.py", line 1052, in run
    while self.step():
  File "/env/torch/_dynamo/symbolic_convert.py", line 962, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/env/torch/_dynamo/symbolic_convert.py", line 659, in wrapper
    return inner_fn(self, inst)
  File "/env/torch/_dynamo/symbolic_convert.py", line 1658, in CALL_FUNCTION
    self.call_function(fn, args, {})
  File "/env/torch/_dynamo/symbolic_convert.py", line 897, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
  File "/env/torch/_dynamo/variables/misc.py", line 1022, in call_function
    return self.obj.call_method(tx, self.name, args, kwargs)
  File "/env/torch/_dynamo/variables/dicts.py", line 566, in call_method
    return super().call_method(tx, name, args, kwargs)
  File "/env/torch/_dynamo/variables/dicts.py", line 396, in call_method
    return super().call_method(tx, name, args, kwargs)
  File "/env/torch/_dynamo/variables/base.py", line 414, in call_method
    unimplemented(f"call_method {self} {name} {args} {kwargs}")
  File "/env/torch/_dynamo/exc.py", line 317, in unimplemented
    raise Unsupported(msg, case_name=case_name)
torch._dynamo.exc.Unsupported: call_method SetVariable() __setitem__ (UserDefinedObjectVariable(GGUFParameter), ConstantVariable(NoneType: None)) {}

from user code:
   File "/env/diffusers/models/transformers/auraflow_transformer_2d.py", line 458, in forward
    temb = self.time_step_embed(timestep).to(dtype=next(self.parameters()).dtype)
  File "/env/torch/nn/modules/module.py", line 2630, in parameters
    for _name, param in self.named_parameters(recurse=recurse):
  File "/env/torch/nn/modules/module.py", line 2657, in named_parameters
    gen = self._named_members(
  File "/env/torch/nn/modules/module.py", line 2604, in _named_members
    memo.add(v)

Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information


You can suppress this exception and fall back to eager by setting:
    import torch._dynamo
    torch._dynamo.config.suppress_errors = True

System Info

• 🤗 Diffusers version: 0.33.0.dev0
• Platform: Linux-6.8.0-52-generic-x86_64-with-glibc2.39
• Running on Google Colab?: No
• Python version: 3.10.16
• PyTorch version (GPU?): 2.6.0+cu124 (True)
• Flax version (CPU?/GPU?/TPU?): not installed (NA)
• Jax version: not installed
• JaxLib version: not installed
• Huggingface_hub version: 0.28.1
• Transformers version: 4.48.2
• Accelerate version: 1.3.0
• PEFT version: not installed
• Bitsandbytes version: not installed
• Safetensors version: 0.5.2
• xFormers version: not installed
• Accelerator: NVIDIA GeForce RTX 4090, 24564 MiB
• Using GPU in script?: no
• Using distributed or parallel set-up in script?: no

Who can help?

@DN6 @hlky @stevhliu

[DN6]

Hi @AstraliteHeart I haven't looked too deeply into it, but I think there are quite a few ops in the GGUFParameter that might not lend well to torch.compile. From the stack trace it seems like it might be related to how we're initializing the GGUFParameter, but I'll have to dig in a bit further to understand if that is actually the case. I don't have the bandwidth at the moment to look into this right now, but I can get to it next week.

[AstraliteHeart]

@DN6 beggars can't be choosers! I appreciate you checking this!
I am learning more on the torch.compile to help figuring this out but it's taking some time unfortunately.

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

[StrongerXi]

I have a fix coming soon, the performance gain from a successful fullgraph=True compilation is pretty significant.

[StrongerXi]

@AstraliteHeart pytorch/pytorch#149483 should resolve the issue, you'll need to install torch at that commit though, and you still need another line to benefit from it:

torch._dynamo.config.traceable_tensor_subclasses.add(diffusers.quantizers.gguf.utils.GGUFParameter)

I'm trying to get it to land asap, after that you can use it by installing torch at main, or nightly, plus the line above.

I'm also trying to see if we can avoid this config line; I'll see how it goes, it's difficult because it means we are opening up a potentially unstable feature.