Fix GGUF KeyError: None when offloading rebuilds a parameter

src/diffusers/quantizers/gguf/utils.py

@@ -548,6 +548,16 @@ def dequantize_gguf_tensor(tensor):
 class GGUFParameter(torch.nn.Parameter):
     def __new__(cls, data, requires_grad=False, quant_type=None):
         data = data if data is not None else torch.empty(0)
+        if quant_type is None:
+            # Offloading rebuilds parameters as `param_cls(new_value, requires_grad=...)` without
+            # forwarding `quant_type` (see `accelerate.utils.set_module_tensor_to_device`), so
+            # inherit it from the tensor being wrapped instead of failing with `KeyError: None`.
+            quant_type = getattr(data, "quant_type", None)
+        if quant_type not in GGML_QUANT_SIZES:
+            raise ValueError(
+                f"`GGUFParameter` expects a valid `quant_type`, but got {quant_type}, and it could not "
+                "be inferred from the tensor being wrapped."
+            )
         self = torch.Tensor._make_subclass(cls, data, requires_grad)
         self.quant_type = quant_type
         block_size, type_size = GGML_QUANT_SIZES[quant_type]

tests/pipelines/testing_utils/quantization.py

@@ -1380,6 +1380,30 @@ def test_pipeline_inference(self):
         max_diff = numpy_cosine_similarity_distance(expected_slice, output_slice)
         assert max_diff < 1e-4
 
+    def test_pipeline_inference_sequential_cpu_offload(self):
+        r"""
+        Sequential CPU offload rebuilds every parameter through `param_cls(new_value, ...)`, which
+        used to drop `GGUFParameter.quant_type` and fail with `KeyError: None`. Like the TorchAO
+        equivalent this only checks that inference runs.
+        """
+        quantization_config = GGUFQuantizationConfig(compute_dtype=self.torch_dtype)
+        transformer = self.model_cls.from_single_file(
+            self.ckpt_path, quantization_config=quantization_config, torch_dtype=self.torch_dtype
+        )
+        pipe = FluxPipeline.from_pretrained(
+            "black-forest-labs/FLUX.1-dev", transformer=transformer, torch_dtype=self.torch_dtype
+        )
+        pipe.enable_sequential_cpu_offload()
+
+        output = pipe(
+            prompt="a cat holding a sign that says hello",
+            num_inference_steps=2,
+            generator=torch.Generator("cpu").manual_seed(0),
+            output_type="np",
+        ).images[0]
+
+        assert output.shape == (1024, 1024, 3)
+
 
 class TestSD35LargeGGUFPipeline(GGUFPipelineTests):
     ckpt_path = "https://huggingface.co/city96/stable-diffusion-3.5-large-gguf/blob/main/sd3.5_large-Q4_0.gguf"

tests/quantization/gguf/test_gguf.py

@@ -83,3 +83,48 @@ def test_cuda_kernels_vs_native(self):
             assert torch.allclose(output_native, output_cuda, 1e-2), (
                 f"GGUF CUDA Kernel Output is different from Native Output for {quant_type}"
             )
+
+
+@is_quantization
+@is_gguf
+@require_accelerate
+@require_gguf_version_greater_or_equal("0.10.0")
+class TestGGUFParameterRewrap:
+    """Offloading rebuilds parameters without forwarding `quant_type`."""
+
+    def _make_param(self):
+        from diffusers.quantizers.gguf.utils import GGML_QUANT_SIZES
+
+        quant_type = gguf.GGMLQuantizationType.Q8_0
+        _, type_size = GGML_QUANT_SIZES[quant_type]
+        data = torch.zeros((4, type_size), dtype=torch.uint8)
+
+        return GGUFParameter(data, quant_type=quant_type), quant_type
+
+    def test_rewrap_without_quant_type_inherits_it(self):
+        param, quant_type = self._make_param()
+        rewrapped = GGUFParameter(param, requires_grad=False)
+
+        assert rewrapped.quant_type == quant_type
+        assert rewrapped.quant_shape == param.quant_shape
+
+    def test_set_module_tensor_to_device_preserves_quant_type(self):
+        from accelerate.utils import set_module_tensor_to_device
+
+        from diffusers.quantizers.gguf.utils import GGUFLinear
+
+        param, quant_type = self._make_param()
+        out_features, in_features = param.quant_shape
+        module = GGUFLinear(in_features, out_features, bias=False, compute_dtype=torch.bfloat16)
+        module.weight = param
+
+        # Passing `value` (or moving across devices) is what makes `accelerate` rebuild the
+        # parameter through `param_cls(new_value, requires_grad=...)`. A same-device move with
+        # no value short-circuits before that branch and would not cover the regression.
+        set_module_tensor_to_device(module, "weight", torch.device("cpu"), value=param)
+
+        assert module.weight.quant_type == quant_type
+
+    def test_untyped_construction_raises(self):
+        with pytest.raises(ValueError):
+            GGUFParameter(torch.zeros((4, 32), dtype=torch.uint8))