refactor pipeline-level quantization tests

Author: sayak@huggingface.co

.github/workflows/nightly_tests.yml

@@ -343,19 +343,19 @@ jobs:
       matrix:
         config:
           - backend: "bitsandbytes"
-            test_location: "bnb"
+            marker: "bitsandbytes"
             additional_deps: ["peft"]
           - backend: "gguf"
-            test_location: "gguf"
+            marker: "gguf"
             additional_deps: ["peft", "kernels"]
           - backend: "torchao"
-            test_location: "torchao"
-            additional_deps: []
+            marker: "torchao"
+            additional_deps: ["mslk"]
           - backend: "optimum_quanto"
-            test_location: "quanto"
+            marker: "quanto"
             additional_deps: []
           - backend: "nvidia_modelopt"
-            test_location: "modelopt"
+            marker: "modelopt"
             additional_deps: []
     runs-on:
       group: aws-g6e-xlarge-plus
@@ -390,9 +390,12 @@ jobs:
           BIG_GPU_MEMORY: 40
         run: |
           pytest -n 1 --max-worker-restart=0 --dist=loadfile \
+            -m "${{ matrix.config.marker }}" \
             --make-reports=tests_${{ matrix.config.backend }}_torch_cuda \
             --report-log=tests_${{ matrix.config.backend }}_torch_cuda.log \
-            tests/quantization/${{ matrix.config.test_location }}
+            tests/models \
+            tests/quantization \
+            tests/pipelines/testing_utils/quantization.py
       - name: Failure short reports
         if: ${{ failure() }}
         run: |
@@ -440,9 +443,10 @@ jobs:
           BIG_GPU_MEMORY: 40
         run: |
           pytest -n 1 --max-worker-restart=0 --dist=loadfile \
+            -k "TestPipelineQuantization" \
             --make-reports=tests_pipeline_level_quant_torch_cuda \
             --report-log=tests_pipeline_level_quant_torch_cuda.log \
-            tests/quantization/test_pipeline_level_quantization.py
+            tests/pipelines/testing_utils/quantization.py
       - name: Failure short reports
         if: ${{ failure() }}
         run: |

tests/models/testing_utils/quantization.py

@@ -16,6 +16,7 @@
 import gc
 
 import pytest
+import safetensors.torch
 import torch
 
 from diffusers import (
@@ -423,6 +424,14 @@ def _test_quantization_device_map(self, config_kwargs):
         assert hasattr(model, "hf_device_map"), "Model should have hf_device_map attribute"
         assert model.hf_device_map is not None, "hf_device_map should not be None"
 
+        map_devices = {torch.device(d).type for d in model.hf_device_map.values()}
+        for name, buffer in model.named_buffers():
+            assert buffer.device.type != "meta", f"Buffer {name} was left on the meta device"
+            if len(map_devices) == 1:
+                assert buffer.device.type == next(iter(map_devices)), (
+                    f"Expected device {next(iter(map_devices))} for buffer {name}, got {buffer.device}"
+                )
+
         inputs = self.get_dummy_inputs()
         output = model(**inputs, return_dict=False)[0]
         assert output is not None, "Model output is None"
@@ -625,15 +634,32 @@ def test_bnb_quantization_memory_footprint(self, config_name):
     def test_bnb_quantization_inference(self, config_name):
         self._test_quantization_inference(BitsAndBytesConfigMixin.BNB_CONFIGS[config_name])
 
-    @pytest.mark.parametrize("config_name", ["4bit_nf4"], ids=["4bit_nf4"])
+    @pytest.mark.parametrize("config_name", ["4bit_nf4", "8bit"], ids=["4bit_nf4", "8bit"])
     def test_bnb_quantization_dtype_assignment(self, config_name):
         self._test_quantization_dtype_assignment(BitsAndBytesConfigMixin.BNB_CONFIGS[config_name])
 
+    def test_bnb_device_assignment(self):
+        """Test that a 4-bit model moves between CPU and accelerator without changing its memory footprint."""
+        model = self._create_quantized_model(BitsAndBytesConfigMixin.BNB_CONFIGS["4bit_nf4"])
+        mem_before = model.get_memory_footprint()
+
+        model.to("cpu")
+        assert model.device.type == "cpu"
+        assert model.get_memory_footprint() == pytest.approx(mem_before)
+
+        model.to(torch_device)
+        assert model.device.type == torch.device(torch_device).type
+        assert model.get_memory_footprint() == pytest.approx(mem_before)
+
     @pytest.mark.parametrize("config_name", ["4bit_nf4"], ids=["4bit_nf4"])
     def test_bnb_quantization_lora_inference(self, config_name):
         self._test_quantization_lora_inference(BitsAndBytesConfigMixin.BNB_CONFIGS[config_name])
 
-    @pytest.mark.parametrize("config_name", ["4bit_nf4"], ids=["4bit_nf4"])
+    @pytest.mark.parametrize(
+        "config_name",
+        list(BitsAndBytesConfigMixin.BNB_CONFIGS.keys()),
+        ids=list(BitsAndBytesConfigMixin.BNB_CONFIGS.keys()),
+    )
     def test_bnb_quantization_serialization(self, config_name, tmp_path):
         self._test_quantization_serialization(BitsAndBytesConfigMixin.BNB_CONFIGS[config_name], tmp_path)
 
@@ -660,15 +686,51 @@ def test_bnb_keep_modules_in_fp32(self):
         self._test_keep_modules_in_fp32(BitsAndBytesConfigMixin.BNB_CONFIGS["4bit_nf4"])
 
     def test_bnb_modules_to_not_convert(self):
-        """Test that modules_to_not_convert parameter works correctly."""
+        """Test that `llm_int8_skip_modules` (the BitsAndBytesConfig module-exclusion option) works correctly."""
         modules_to_exclude = getattr(self, "modules_to_not_convert_for_test", None)
         if modules_to_exclude is None:
             pytest.skip("modules_to_not_convert_for_test not defined for this model")
 
-        self._test_quantization_modules_to_not_convert(
-            BitsAndBytesConfigMixin.BNB_CONFIGS["4bit_nf4"], modules_to_exclude
+        config_kwargs = {**BitsAndBytesConfigMixin.BNB_CONFIGS["8bit"], "llm_int8_skip_modules": modules_to_exclude}
+        model = self._create_quantized_model(config_kwargs)
+
+        found_excluded = False
+        for name, module in model.named_modules():
+            if isinstance(module, torch.nn.Linear):
+                if any(excluded in name for excluded in modules_to_exclude):
+                    found_excluded = True
+                    assert module.weight.dtype != torch.int8, f"Module {name} should not be quantized"
+                else:
+                    assert isinstance(module, bnb.nn.Linear8bitLt), f"Module {name} should be quantized"
+                    assert module.weight.dtype == torch.int8, f"Module {name} weight should be int8"
+
+        assert found_excluded, f"No linear layers found in excluded modules: {modules_to_exclude}"
+
+    @pytest.mark.parametrize("config_name", ["8bit"], ids=["8bit"])
+    def test_bnb_quantization_sharded_serialization(self, config_name, tmp_path):
+        self._test_quantization_serialization(
+            BitsAndBytesConfigMixin.BNB_CONFIGS[config_name], tmp_path, max_shard_size="16KB"
         )
 
+    def test_bnb_errors_loading_incorrect_state_dict(self, tmp_path):
+        """Test that loading a checkpoint with a corrupted quantized weight raises a helpful error."""
+        model = self._create_quantized_model(BitsAndBytesConfigMixin.BNB_CONFIGS["4bit_nf4"])
+        model.save_pretrained(str(tmp_path))
+        del model
+        gc.collect()
+        backend_empty_cache(torch_device)
+
+        weights_file = tmp_path / "diffusion_pytorch_model.safetensors"
+        state_dict = safetensors.torch.load_file(str(weights_file))
+        key_to_target = next(k for k in state_dict if k.endswith(".weight") and state_dict[k].dtype == torch.uint8)
+        corrupted_param = torch.randn(state_dict[key_to_target].shape[0] - 1, 1)
+        state_dict[key_to_target] = bnb.nn.Params4bit(corrupted_param, requires_grad=False)
+        safetensors.torch.save_file(state_dict, str(weights_file))
+
+        with pytest.raises(ValueError) as err_context:
+            _ = self.model_class.from_pretrained(str(tmp_path))
+        assert key_to_target in str(err_context.value)
+
     @pytest.mark.parametrize("config_name", ["4bit_nf4", "8bit"], ids=["4bit_nf4", "8bit"])
     def test_bnb_device_map(self, config_name):
         """Test that device_map='auto' works correctly with quantization."""
@@ -678,9 +740,10 @@ def test_bnb_dequantize(self):
         """Test that dequantize() works correctly."""
         self._test_dequantize(BitsAndBytesConfigMixin.BNB_CONFIGS["4bit_nf4"])
 
-    def test_bnb_training(self):
+    @pytest.mark.parametrize("config_name", ["4bit_nf4", "8bit"], ids=["4bit_nf4", "8bit"])
+    def test_bnb_training(self, config_name):
         """Test that quantized models can be used for training with adapters."""
-        self._test_quantization_training(BitsAndBytesConfigMixin.BNB_CONFIGS["4bit_nf4"])
+        self._test_quantization_training(BitsAndBytesConfigMixin.BNB_CONFIGS[config_name])
 
     @pytest.mark.parametrize(
         "config_name",
@@ -812,6 +875,10 @@ def test_quanto_quantization_inference(self, weight_type_name):
     def test_quanto_quantized_layers(self, weight_type_name):
         self._test_quantized_layers(QuantoConfigMixin.QUANTO_WEIGHT_TYPES[weight_type_name])
 
+    @pytest.mark.parametrize("weight_type_name", ["int8"], ids=["int8"])
+    def test_quanto_quantization_dtype_assignment(self, weight_type_name):
+        self._test_quantization_dtype_assignment(QuantoConfigMixin.QUANTO_WEIGHT_TYPES[weight_type_name])
+
     @pytest.mark.parametrize("weight_type_name", ["int8"], ids=["int8"])
     def test_quanto_quantization_lora_inference(self, weight_type_name):
         self._test_quantization_lora_inference(QuantoConfigMixin.QUANTO_WEIGHT_TYPES[weight_type_name])
@@ -1016,6 +1083,47 @@ def test_torchao_device_map(self):
         """Test that device_map='auto' works correctly with quantization."""
         self._test_quantization_device_map(TorchAoConfigMixin.TORCHAO_QUANT_TYPES["int8wo"])
 
+    @torch.no_grad()
+    def test_torchao_cpu_disk_offload_device_map(self, tmp_path):
+        """Test custom device maps with cpu/disk offload: offloaded modules stay unquantized, inference works."""
+        from torchao.utils import TorchAOBaseTensor
+
+        model = self._create_quantized_model(TorchAoConfigMixin.TORCHAO_QUANT_TYPES["int8wo"])
+
+        # Offload the first two linear-bearing top-level modules to cpu and disk, keep the rest on the accelerator.
+        device_map = {}
+        offload_targets = []
+        for name, child in model.named_children():
+            if len(offload_targets) < 2 and any(isinstance(m, torch.nn.Linear) for m in child.modules()):
+                device_map[name] = "disk" if offload_targets else "cpu"
+                offload_targets.append(name)
+            else:
+                device_map[name] = str(torch_device)
+        del model
+        gc.collect()
+        backend_empty_cache(torch_device)
+        if len(offload_targets) < 2:
+            pytest.skip("Model does not have enough linear-bearing top-level modules for offload testing")
+
+        model = self._create_quantized_model(
+            TorchAoConfigMixin.TORCHAO_QUANT_TYPES["int8wo"], device_map=device_map, offload_folder=str(tmp_path)
+        )
+
+        # Weights offloaded to cpu/disk are not quantized, only the weights on the accelerator are.
+        for name, module in model.named_modules():
+            if isinstance(module, torch.nn.Linear):
+                if name.split(".")[0] in offload_targets:
+                    assert not isinstance(module.weight, TorchAOBaseTensor), (
+                        f"Offloaded module {name} should not be quantized"
+                    )
+                else:
+                    assert isinstance(module.weight, TorchAOBaseTensor), f"Module {name} should be quantized"
+
+        inputs = self.get_dummy_inputs()
+        output = model(**inputs, return_dict=False)[0]
+        assert output is not None, "Model output is None"
+        assert not torch.isnan(output).any(), "Model output contains NaN"
+
     @pytest.mark.parametrize(
         "quant_type",
         [
@@ -1103,6 +1211,38 @@ class GGUFTesterMixin(GGUFConfigMixin, QuantizationTesterMixin):
     def test_gguf_quantization_inference(self):
         self._test_quantization_inference({"compute_dtype": torch.bfloat16})
 
+    def test_gguf_quantized_layers(self):
+        compute_dtype = getattr(self, "torch_dtype", torch.bfloat16)
+        model = self._create_quantized_model({"compute_dtype": compute_dtype})
+
+        num_quantized = 0
+        for name, module in model.named_modules():
+            if isinstance(module, torch.nn.Linear) and hasattr(module.weight, "quant_type"):
+                self._verify_if_layer_quantized(name, module)
+                if module.bias is not None:
+                    assert module.bias.dtype == compute_dtype, f"{name} bias should be {compute_dtype}"
+                num_quantized += 1
+
+        assert num_quantized > 0, "No quantized linear layers found in model"
+
+    @torch.no_grad()
+    def test_gguf_memory_usage(self):
+        expected_gb = getattr(self, "expected_memory_use_in_gb", None)
+        if expected_gb is None:
+            pytest.skip("expected_memory_use_in_gb not defined for this model")
+
+        compute_dtype = getattr(self, "torch_dtype", torch.bfloat16)
+        model = self._create_quantized_model({"compute_dtype": compute_dtype})
+        model.to(torch_device)
+        assert (model.get_memory_footprint() / 1024**3) < expected_gb
+
+        inputs = self.get_dummy_inputs()
+        backend_reset_peak_memory_stats(torch_device)
+        backend_empty_cache(torch_device)
+        model(**inputs)
+        max_memory = backend_max_memory_allocated(torch_device)
+        assert (max_memory / 1024**3) < expected_gb
+
     def test_gguf_keep_modules_in_fp32(self):
         if not hasattr(self.model_class, "_keep_in_fp32_modules"):
             pytest.skip(f"{self.model_class.__name__} does not have _keep_in_fp32_modules")

tests/models/transformers/test_models_transformer_flux.py

@@ -19,7 +19,7 @@
 import pytest
 import torch
 
-from diffusers import BitsAndBytesConfig, FluxTransformer2DModel
+from diffusers import BitsAndBytesConfig, FluxTransformer2DModel, GGUFQuantizationConfig
 from diffusers.models.embeddings import ImageProjection
 from diffusers.models.transformers.transformer_flux import FluxIPAdapterAttnProcessor
 from diffusers.utils.torch_utils import randn_tensor
@@ -356,6 +356,8 @@ def torch_dtype(self):
 class TestFluxTransformerBitsAndBytes(FluxTransformerTesterConfig, BitsAndBytesTesterMixin):
     """BitsAndBytes quantization tests for Flux Transformer."""
 
+    modules_to_not_convert_for_test = ["proj_out"]
+
     @property
     def torch_dtype(self):
         return torch.float16
@@ -384,6 +386,8 @@ def torch_dtype(self):
 
 
 class TestFluxTransformerGGUF(FluxTransformerTesterConfig, GGUFTesterMixin):
+    expected_memory_use_in_gb = 5
+
     @property
     def gguf_filename(self):
         return "https://huggingface.co/city96/FLUX.1-dev-gguf/blob/main/flux1-dev-Q2_K.gguf"
@@ -410,6 +414,17 @@ def get_dummy_inputs(self):
             "guidance": torch.tensor([3.5]).to(torch_device, self.torch_dtype),
         }
 
+    @torch.no_grad()
+    def test_loading_gguf_diffusers_format(self):
+        model = self.model_class.from_single_file(
+            "https://huggingface.co/sayakpaul/flux-diffusers-gguf/blob/main/model-Q4_0.gguf",
+            subfolder="transformer",
+            quantization_config=GGUFQuantizationConfig(compute_dtype=self.torch_dtype),
+            config="black-forest-labs/FLUX.1-dev",
+        )
+        model.to(torch_device)
+        model(**self.get_dummy_inputs())
+
 
 class TestFluxTransformerQuantoCompile(FluxTransformerTesterConfig, QuantoCompileTesterMixin):
     """Quanto + compile tests for Flux Transformer."""

tests/models/transformers/test_models_transformer_wan_animate.py

@@ -239,6 +239,8 @@ def get_dummy_inputs(self):
 class TestWanAnimateTransformer3DGGUF(WanAnimateTransformer3DTesterConfig, GGUFTesterMixin):
     """GGUF quantization tests for Wan Animate Transformer 3D."""
 
+    expected_memory_use_in_gb = 9
+
     @property
     def gguf_filename(self):
         return "https://huggingface.co/QuantStack/Wan2.2-Animate-14B-GGUF/blob/main/Wan2.2-Animate-14B-Q2_K.gguf"

tests/models/transformers/test_models_transformer_wan_vace.py

@@ -196,6 +196,8 @@ def get_dummy_inputs(self):
 class TestWanVACETransformer3DGGUF(WanVACETransformer3DTesterConfig, GGUFTesterMixin):
     """GGUF quantization tests for Wan VACE Transformer 3D."""
 
+    expected_memory_use_in_gb = 9
+
     @property
     def gguf_filename(self):
         return "https://huggingface.co/QuantStack/Wan2.1_14B_VACE-GGUF/blob/main/Wan2.1_14B_VACE-Q3_K_S.gguf"