GGUF: enable_sequential_cpu_offload() raises KeyError: None — any GGUF model larger than VRAM is unrunnable

[EigenAx2Pi]

Describe the bug

GGUFParameter.__new__ accepts quant_type=None as its default and then immediately performs an
unguarded dict lookup on it:

# src/diffusers/quantizers/gguf/utils.py
def __new__(cls, data, requires_grad=False, quant_type=None):
    data = data if data is not None else torch.empty(0)
    self = torch.Tensor._make_subclass(cls, data, requires_grad)
    self.quant_type = quant_type
    block_size, type_size = GGML_QUANT_SIZES[quant_type]   # <-- KeyError when quant_type is None
    self.quant_shape = _quant_shape_from_byte_shape(self.shape, type_size, block_size)
    return self

accelerate's offload path (set_module_tensor_to_device) re-wraps parameters as
param_cls(new_value, requires_grad=old_value.requires_grad) — without forwarding
quant_type. Every such re-wrap therefore raises KeyError: None.

Impact. This makes enable_sequential_cpu_offload() unusable with any GGUF-quantised
transformer. Because enable_model_cpu_offload() moves the whole transformer to the GPU at once
(and fails with a driver-level OOM when it does not fit), the practical effect is that no GGUF
model larger than available VRAM can be run at all — precisely the case GGUF quantisation exists
to serve. On an 8 GB card this ruled out FLUX.1-schnell Q8_0 (12.7 GB) and Chroma1-HD Q8_0 (9.7 GB).

The error is also misleading: KeyError: None surfacing from inside accelerate reads as a corrupt
or unsupported model file. We wrote off three different models as "broken quants" before running a
known-good file under seq offload and finding it failed identically.

Reproduction

Minimal — no model download required:

import torch
from diffusers.quantizers.gguf.utils import GGUFParameter

p = GGUFParameter(torch.zeros(32, 32), quant_type=8)  # Q8_0
GGUFParameter(p, requires_grad=False)                 # KeyError: None

End-to-end:

import torch
from diffusers import FluxPipeline, FluxTransformer2DModel, GGUFQuantizationConfig

tr = FluxTransformer2DModel.from_single_file(
    "https://huggingface.co/city96/FLUX.1-schnell-gguf/blob/main/flux1-schnell-Q4_K_S.gguf",
    quantization_config=GGUFQuantizationConfig(compute_dtype=torch.bfloat16),
    torch_dtype=torch.bfloat16,
    config="black-forest-labs/FLUX.1-schnell", subfolder="transformer",
)
pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell",
                                    transformer=tr, torch_dtype=torch.bfloat16)
pipe.enable_sequential_cpu_offload()          # works with enable_model_cpu_offload()
pipe("a cat", num_inference_steps=1)          # KeyError: None

Suggested fix

Inherit quant_type from the incoming tensor when it is not passed explicitly, and tolerate a
genuinely-absent type rather than raising. quant_shape is assigned but never read elsewhere in
diffusers, so leaving it None in that case is safe.

if quant_type is None:
    quant_type = getattr(data, "quant_type", None)
self = torch.Tensor._make_subclass(cls, data, requires_grad)
self.quant_type = quant_type
if quant_type is None:
    self.quant_shape = None
else:
    block_size, type_size = GGML_QUANT_SIZES[quant_type]
    self.quant_shape = _quant_shape_from_byte_shape(self.shape, type_size, block_size)

_extract_quant_type should probably also inspect kwargs, not only positional args.

Result after patching

enable_sequential_cpu_offload() works, and peak VRAM for FLUX.1-schnell Q4_K_S drops from
7.4 GB to 1.81 GB (26.9s vs 18s per 1024x1024 image, 4 steps, RTX 5070 Laptop 8 GB).
Models that previously could not be loaded at all now run.

System info

• diffusers 0.40.0, torch 2.11.0+cu128, accelerate (bundled), gguf 0.19.0
• Python 3.12, Linux (WSL2), RTX 5070 Laptop 8 GB (sm_120)

[TANGBUDU]

@EigenAx2Pi are you already preparing a PR for this? If not, I'd like to take it.

One note on the proposed fix: quant_shape is still read outside gguf/utils.py. FluxLoraLoaderMixin._calculate_module_shape() returns weight.quant_shape for a GGUFParameter, and its callers expect a real shape, so quant_shape=None would just move the failure into the Flux LoRA path. Same for dequantize_gguf_tensor(): it only checks that quant_type exists, not that it's valid, so a GGUFParameter with quant_type=None still fails later.

Inheriting quant_type from data looks right to me for the reported Accelerate re-wrap path. I'd rather keep genuinely untyped construction invalid though, instead of ending up with a GGUFParameter that carries no quantization metadata (that construction already raises today, so nothing that works now would start failing).

I'd also add a test for the Accelerate re-wrap itself, plus one for GGUF sequential CPU offload: among the quantization tests, sequential offload is only covered on the TorchAO side, and the GGUF pipeline tests use model CPU offload.

@sayakpaul, if that scope sounds right, I'm happy to work on a PR.

[DN6]

Thanks for flagging @EigenAx2Pi. Assigning this issue to you. Btw have you tried running it with leaf level group offloading? Does the issue persist?

[EigenAx2Pi]

@DN6 thanks — I measured it, and leaf-level group offloading is not affected. Same transformer, same box, with my local fix reverted so it couldn't mask anything:

offload mode	result	peak VRAM	time
enable_sequential_cpu_offload()	KeyError: None	—	—
group offload, leaf_level	works	0.77 GiB	8.3–11.5 s
group offload, leaf_level, use_stream=True	works	0.77 GiB	3.0–3.7 s
enable_sequential_cpu_offload() with the fix	works	0.61 GiB	5.4–6.5 s

FLUX.1-schnell Q4_K_S, 512×512, 1 step, RTX 5070 Laptop 8 GB. Text encoders nulled and embeds passed directly, so only the transformer offload path is exercised; the VAE stays resident in the group-offload rows, which is where most of the 0.77 vs 0.61 GiB difference comes from.

It doesn't reproduce for a mechanical reason: group offloading moves tensors in place (param.data = param.data.to(...) in hooks/group_offloading.py), so GGUFParameter.__new__ is never called a second time and quant_type survives. enable_sequential_cpu_offload() goes through accelerate.utils.set_module_tensor_to_device, which rebuilds the parameter — param_cls(new_value, requires_grad=old_value.requires_grad) — and drops quant_type. Any path that reconstructs the parameter hits this; any path that mutates .data does not.

So group offload with streams is a good workaround today, and on this box it's actually faster than sequential offload. But enable_sequential_cpu_offload() is the first thing the low-VRAM docs point at, it's broken for every GGUF checkpoint, and the symptom — a bare KeyError: None surfacing from inside accelerate — reads like a corrupt quant file rather than a library bug. I think it's worth fixing regardless of the workaround.

---

@TANGBUDU you're right on both counts, thank you — I've dropped the quant_shape = None branch. _calculate_module_shape() does read weight.quant_shape for a GGUFParameter in loaders/lora_pipeline.py, and dequantize_gguf_tensor() only checks that quant_type exists, not that it's valid, so a None-typed parameter would have relocated the failure rather than removed it. My "never read elsewhere" claim in the issue was simply wrong.

What I'm taking forward is inheritance only, plus an explicit error where there's genuinely nothing to inherit — which already raises today, so nothing that currently works starts failing:

if quant_type is None:
    quant_type = getattr(data, "quant_type", None)
if quant_type not in GGML_QUANT_SIZES:
    raise ValueError(...)

Verified end-to-end with that exact body: sequential offload runs, 0.61 GiB peak.

On tests — agreed, and it still holds after the test reorg: test_sequential_cpu_offload exists only on TestTorchAo, and all four GGUF pipeline tests (TestFluxGGUFPipeline, TestSD35Large…, TestSD35Medium…, TestAuraFlow…) use enable_model_cpu_offload(). I'm adding two backend-level re-wrap tests — the direct GGUFParameter(param) re-wrap and the set_module_tensor_to_device path, both of which fail on main with KeyError: None and pass with the fix — plus a GGUF sequential-offload pipeline test. Worth flagging that the pipeline one is nightly/big-accelerator, so I can't run it on an 8 GB card; CI will be its first real execution.

One trap for whoever reviews: a naive set_module_tensor_to_device test passes on main, because a same-device move short-circuits before the re-wrap branch. It only reproduces on an actual device change or an explicit value=.

@DN6 assigned this to me so I'll open the PR — @TANGBUDU I'd genuinely value your review on it, and I'll credit the quant_shape catch in the description.

---

Env (for the needs-env-info label)

- 🤗 Diffusers version: 0.40.0
- Platform: Linux-6.6.87.2-microsoft-standard-WSL2-x86_64-with-glibc2.39
- Running on Google Colab?: No
- Python version: 3.12.3
- PyTorch version (GPU?): 2.11.0+cu128 (True)
- Huggingface_hub version: 1.29.0
- Transformers version: 5.16.1
- Accelerate version: 1.14.0
- gguf version: 0.19.0
- Safetensors version: 0.8.0
- Accelerator: NVIDIA GeForce RTX 5070 Laptop GPU, 8151 MiB
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: No