Fix GGUF KeyError: None when offloading rebuilds a parameter

[EigenAx2Pi]

Fixes #14691.

What was broken

GGUFParameter.__new__ defaults quant_type to None and then looks it up in GGML_QUANT_SIZES without a guard. Offloading rebuilds parameters as param_cls(new_value, requires_grad=old_value.requires_grad) — accelerate.utils.set_module_tensor_to_device — without forwarding quant_type, so every rebuild raised KeyError: None.

The effect: enable_sequential_cpu_offload() was unusable with any GGUF-quantised transformer. Since enable_model_cpu_offload() moves the whole transformer onto the accelerator at once, no GGUF checkpoint larger than available VRAM could be run at all — precisely the case GGUF quantisation exists to serve. On an 8 GB card that ruled out FLUX.1-schnell Q8_0 (12.7 GB) and Chroma1-HD Q8_0 (9.7 GB).

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

The fix

Inherit quant_type from the tensor being wrapped when it isn't passed, and raise an explicit error where there is nothing to inherit. That second path already raised (GGML_QUANT_SIZES[None]), so nothing that works today starts failing — it only stops the failure from surfacing as a bare KeyError: None.

The original suggestion in the issue was to leave quant_shape = None instead of raising. @TANGBUDU correctly pointed out that this is wrong: FluxLoraLoaderMixin._calculate_module_shape() does read weight.quant_shape for a GGUFParameter, 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. Credit to them for the catch; this PR takes the inheritance-only route because of it.

Tests

Two backend-level regression tests plus one pipeline-level test:

test	on main	with this PR
test_rewrap_without_quant_type_inherits_it	fails, KeyError: None	passes
test_set_module_tensor_to_device_preserves_quant_type	fails, KeyError: None	passes
test_untyped_construction_raises	fails (raises KeyError)	passes

Verified in both directions locally (pytest tests/quantization/gguf/test_gguf.py -k Rewrap), reverting the utils.py change to confirm each one actually fails without it.

One trap worth flagging for review: a naive set_module_tensor_to_device test passes on main, because a same-device move with no value short-circuits before the rebuild branch. It only reproduces on an actual device change or an explicit value=.

test_pipeline_inference_sequential_cpu_offload closes the coverage gap @TANGBUDU identified — test_sequential_cpu_offload currently exists only on TestTorchAo, and all four GGUF pipeline tests use enable_model_cpu_offload(). It follows the TorchAO precedent of asserting only that inference runs. It is nightly + big-accelerator, so I could not run it on an 8 GB card — CI will be its first real execution.

On group offloading

@DN6 asked whether the issue persists under leaf-level group offloading. It does not, and the reason is mechanical: group offloading moves tensors in place (param.data = param.data.to(...) in hooks/group_offloading.py), so GGUFParameter.__new__ is never re-entered. Only paths that reconstruct the parameter are affected.

Measured on FLUX.1-schnell Q4_K_S, 512×512, 1 step, RTX 5070 Laptop 8 GB (text encoders nulled so only the transformer offload path is exercised):

offload mode	result	peak VRAM	time
enable_sequential_cpu_offload(), main	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(), this PR	works	0.61 GiB	5.4–6.5 s

So group offloading is a real workaround today. enable_sequential_cpu_offload() is still the API the low-VRAM docs point at first, though, and it is broken for every GGUF checkpoint.

Env

diffusers main, torch 2.11.0+cu128, accelerate 1.14.0, gguf 0.19.0, Python 3.12, Linux (WSL2), RTX 5070 Laptop 8 GB (sm_120).

cc @DN6 @TANGBUDU — review very welcome, this is my first contribution here.