GGUF: torch.compile cannot trace sets

[dxqb]

torch.compile fails to trace

s = { ... }
if x in s:

it can trace lists

Who can review?

@Isotr0py

[yiyixuxu]

we use set operations on them like here

diffusers/src/diffusers/quantizers/gguf/utils.py

Line 74 in df8dd77

DEQUANT_TYPES = STANDARD_QUANT_TYPES | KQUANT_TYPES | IMATRIX_QUANT_TYPES

would this work?

[dxqb]

you are right, this would not work and another solution is required

[yiyixuxu]

thanks for the PR! could you provide a testing script?

[dxqb]

thanks for the PR! could you provide a testing script?

import torch
import gguf

UNQUANTIZED_TYPES = {gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16, gguf.GGMLQuantizationType.BF16}

def f(x):
    if x in UNQUANTIZED_TYPES:
        pass

f_compiled = torch.compile(f, fullgraph=True)
f_compiled(gguf.GGMLQuantizationType.F32)

torch._dynamo.exc.Unsupported: Unsupported method call
  Explanation: Dynamo does not know how to trace method `__contains__` of class `set`
  Hint: Avoid calling `set.__contains__` in your code.

[DN6]

@dxqb Could you provide an example where compile is failing with GGUF? I'm able to run this snippet without any issues?

import torch
from diffusers import FluxPipeline, FluxTransformer2DModel, GGUFQuantizationConfig

model_id = "black-forest-labs/FLUX.1-dev"
ckpt_path = (
    "https://huggingface.co/city96/FLUX.1-dev-gguf/blob/main/flux1-dev-Q2_K.gguf"
)
transformer = FluxTransformer2DModel.from_single_file(
    ckpt_path,
    quantization_config=GGUFQuantizationConfig(compute_dtype=torch.bfloat16),
    torch_dtype=torch.bfloat16,
)
transformer = torch.compile(transformer, fullgraph=True)

pipe = FluxPipeline.from_pretrained(
    model_id, transformer=transformer, torch_dtype=torch.bfloat16
)
pipe.to("cuda")
images = pipe("A cat holding a sign that says hello").images[0]
images.save("flux-gguf-compile.png")

[dxqb]

@dxqb Could you provide an example where compile is failing with GGUF? I'm able to run this snippet without any issues?

the untracable line is here:

diffusers/src/diffusers/quantizers/gguf/utils.py

Line 81 in 1a1634d

if qweight_type in UNQUANTIZED_TYPES:

That code path is only reached if Isotr0py's kernels are used so it depends on several conditions, here:

diffusers/src/diffusers/quantizers/gguf/utils.py

Line 33 in 1a1634d

os.getenv("DIFFUSERS_GGUF_CUDA_KERNELS", "false").lower() in ["1", "true", "yes"]

When all these conditions are fulfilled, you should see it in your example. I've posted above a simpler test case that contains the same line so you can reproduce it in any case.