# coding=utf-8 # Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import gc import re import tempfile import pytest import safetensors.torch import torch from diffusers import ( AutoRoundConfig, BFLQuantizationConfig, BitsAndBytesConfig, GGUFQuantizationConfig, NunchakuLiteQuantizationConfig, NVIDIAModelOptConfig, SDNQConfig, TorchAoConfig, ) from diffusers.loaders.single_file_model import SINGLE_FILE_LOADABLE_CLASSES, _get_single_file_loadable_mapping_class from diffusers.utils.import_utils import ( is_bitsandbytes_available, is_flashinfer_available, is_gguf_available, is_kernels_available, is_nvidia_modelopt_available, is_peft_available, is_torchao_available, ) from ...testing_utils import ( assert_tensors_close, backend_empty_cache, is_autoround, is_bfl, is_bitsandbytes, is_gguf, is_modelopt, is_quantization, is_sdnq, is_torch_compile, is_torchao, require_accelerate, require_accelerator, require_auto_round_version_greater_or_equal, require_bitsandbytes_version_greater, require_gguf_version_greater_or_equal, require_modelopt_version_greater_or_equal, require_sdnq, require_torchao_version_greater_or_equal, torch_device, ) from .single_file import fetch_checkpoint_metadata if is_peft_available(): from peft import LoraConfig from diffusers.loaders.peft import PeftAdapterMixin if is_nvidia_modelopt_available(): import modelopt.torch.quantization as mtq if is_bitsandbytes_available(): import bitsandbytes as bnb if is_gguf_available(): pass if is_torchao_available(): import torchao.quantization as _torchao_quantization class LoRALayer(torch.nn.Module): """Wraps a linear layer with LoRA-like adapter - Used for testing purposes only. Taken from https://github.com/huggingface/transformers/blob/566302686a71de14125717dea9a6a45b24d42b37/tests/quantization/bnb/test_4bit.py#L62C5-L78C77 """ def __init__(self, module: torch.nn.Module, rank: int): super().__init__() self.module = module self.adapter = torch.nn.Sequential( torch.nn.Linear(module.in_features, rank, bias=False), torch.nn.Linear(rank, module.out_features, bias=False), ) small_std = (2.0 / (5 * min(module.in_features, module.out_features))) ** 0.5 torch.nn.init.normal_(self.adapter[0].weight, std=small_std) torch.nn.init.zeros_(self.adapter[1].weight) self.adapter.to(module.weight.device) def forward(self, input, *args, **kwargs): return self.module(input, *args, **kwargs) + self.adapter(input) @is_quantization @require_accelerator class QuantizationTesterMixin: """ Base mixin class providing common test implementations for quantization testing. Backend-specific mixins should: 1. Implement _create_quantized_model(config_kwargs) 2. Implement _verify_if_layer_quantized(name, module, config_kwargs) 3. Define their config dict (e.g., BNB_CONFIGS, TORCHAO_CONFIGS, etc.) 4. Use @pytest.mark.parametrize to create tests that call the common test methods below Expected class attributes: - model_class: The model class to test - pretrained_model_name_or_path: Hub repository ID for the pretrained model - pretrained_model_kwargs: (Optional) Dict of kwargs to pass to from_pretrained (e.g., {"subfolder": "transformer"}) Expected methods in test classes: - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass """ def setup_method(self): gc.collect() backend_empty_cache(torch_device) def teardown_method(self): gc.collect() backend_empty_cache(torch_device) def _create_quantized_model(self, config_kwargs, **extra_kwargs): """ Create a quantized model with the given config kwargs. Args: config_kwargs: Quantization config parameters **extra_kwargs: Additional kwargs to pass to from_pretrained (e.g., device_map, offload_folder) """ raise NotImplementedError("Subclass must implement _create_quantized_model") def _verify_if_layer_quantized(self, name, module, config_kwargs): raise NotImplementedError("Subclass must implement _verify_if_layer_quantized") def _is_module_quantized(self, module, config_kwargs=None): """ Check if a module is quantized. Returns True if quantized, False otherwise. Default implementation tries _verify_if_layer_quantized and catches exceptions. Backends whose verifier depends on the quantization config (e.g. bnb's 4-bit/8-bit split) need config_kwargs. Subclasses can override for more efficient checking. """ try: self._verify_if_layer_quantized("", module, config_kwargs or {}) return True except (AssertionError, AttributeError): return False def _load_unquantized_model(self): kwargs = getattr(self, "pretrained_model_kwargs", {}) return self.model_class.from_pretrained(self.pretrained_model_name_or_path, **kwargs) def _test_quantization_num_parameters(self, config_kwargs): model = self._load_unquantized_model() num_params = model.num_parameters() model_quantized = self._create_quantized_model(config_kwargs) num_params_quantized = model_quantized.num_parameters() assert num_params == num_params_quantized, ( f"Parameter count mismatch: unquantized={num_params}, quantized={num_params_quantized}" ) def _test_quantization_memory_footprint(self, config_kwargs, expected_memory_reduction=1.2): model = self._load_unquantized_model() mem = model.get_memory_footprint() model_quantized = self._create_quantized_model(config_kwargs) mem_quantized = model_quantized.get_memory_footprint() ratio = mem / mem_quantized assert ratio >= expected_memory_reduction, ( f"Memory ratio {ratio:.2f} is less than expected ({expected_memory_reduction}x). unquantized={mem}, quantized={mem_quantized}" ) @torch.no_grad() def _test_quantization_inference(self, config_kwargs): model_quantized = self._create_quantized_model(config_kwargs) model_quantized.to(torch_device) inputs = self.get_dummy_inputs() output = model_quantized(**inputs, return_dict=False)[0] assert output is not None, "Model output is None" assert not torch.isnan(output).any(), "Model output contains NaN" def _test_quantization_dtype_assignment(self, config_kwargs): model = self._create_quantized_model(config_kwargs) with pytest.raises(ValueError): model.to(torch.float16) with pytest.raises(ValueError): device_0 = f"{torch_device}:0" model.to(device=device_0, dtype=torch.float16) with pytest.raises(ValueError): model.float() with pytest.raises(ValueError): model.half() model.to(torch_device) @torch.no_grad() def _test_quantization_lora_inference(self, config_kwargs): if not is_peft_available(): pytest.skip("peft is not available") if not issubclass(self.model_class, PeftAdapterMixin): pytest.skip(f"PEFT is not supported for this model ({self.model_class.__name__})") model = self._create_quantized_model(config_kwargs) lora_config = LoraConfig( r=4, lora_alpha=4, target_modules=["to_q", "to_k", "to_v", "to_out.0"], init_lora_weights=False, ) model.add_adapter(lora_config) # Move LoRA adapter weights to device (they default to CPU) model.to(torch_device) inputs = self.get_dummy_inputs() output = model(**inputs, return_dict=False)[0] assert output is not None, "Model output is None with LoRA" assert not torch.isnan(output).any(), "Model output contains NaN with LoRA" # Backends opt into the sharded-serialization test by setting this to a quantization config dict. sharded_serialization_config = None def test_quantization_sharded_serialization(self, tmp_path): if self.sharded_serialization_config is None: pytest.skip("sharded_serialization_config not defined for this backend") self._test_quantization_serialization(self.sharded_serialization_config, tmp_path, max_shard_size="auto") @torch.no_grad() def _test_quantization_serialization(self, config_kwargs, tmp_path, max_shard_size=None): """ Test that a quantized model can be saved and reloaded without changing its outputs. Args: config_kwargs: Quantization config parameters tmp_path: Directory the model is serialized into max_shard_size: When set, the checkpoint is sharded and the shard/index files are checked. "auto" derives a size from the model footprint that yields a handful of shards; sizes far below the largest tensor can split a weight from its quantization components (e.g. bnb's SCB), which the shard-by-shard loader does not support. """ model = self._create_quantized_model(config_kwargs) model.to(torch_device) inputs = self.get_dummy_inputs() expected_output = model(**inputs, return_dict=False)[0].detach().cpu() if max_shard_size == "auto": max_shard_size = max(int(model.get_memory_footprint() // 2), 1) save_kwargs = {"safe_serialization": True} if max_shard_size is not None: save_kwargs["max_shard_size"] = max_shard_size model.save_pretrained(str(tmp_path), **save_kwargs) del model gc.collect() backend_empty_cache(torch_device) if max_shard_size is not None: assert len(list(tmp_path.glob("*.safetensors"))) > 1, "Expected a sharded safe-serialization checkpoint." assert any(path.name.endswith(".index.json") for path in tmp_path.iterdir()), ( "Expected an index file for sharded safe checkpoint." ) model_loaded = self.model_class.from_pretrained(str(tmp_path), device_map=str(torch_device)) output = model_loaded(**inputs, return_dict=False)[0].detach().cpu() assert_tensors_close(output, expected_output, rtol=1e-3, atol=1e-3) def _test_quantization_config_serialization(self, config_kwargs): """ Test that the quantization config attached to a quantized model is serializable. Args: config_kwargs: Quantization config parameters """ model = self._create_quantized_model(config_kwargs) assert "quantization_config" in model.config, "Missing quantization_config" _ = model.config["quantization_config"].to_dict() _ = model.config["quantization_config"].to_diff_dict() _ = model.config["quantization_config"].to_json_string() def _test_original_dtype(self, config_kwargs): """ Test that the dtype the model had before quantization is recorded on its config. Args: config_kwargs: Quantization config parameters """ model = self._create_quantized_model(config_kwargs) assert "_pre_quantization_dtype" in model.config, "Missing _pre_quantization_dtype" assert model.config["_pre_quantization_dtype"] in [ torch.float16, torch.float32, torch.bfloat16, ], f"Unexpected dtype: {model.config['_pre_quantization_dtype']}" def _test_quantized_layers(self, config_kwargs): model_fp = self._load_unquantized_model() num_linear_layers = sum(1 for module in model_fp.modules() if isinstance(module, torch.nn.Linear)) model_quantized = self._create_quantized_model(config_kwargs) num_fp32_modules = 0 if hasattr(model_quantized, "_keep_in_fp32_modules") and model_quantized._keep_in_fp32_modules: for name, module in model_quantized.named_modules(): if isinstance(module, torch.nn.Linear): if any(fp32_name in name for fp32_name in model_quantized._keep_in_fp32_modules): num_fp32_modules += 1 expected_quantized_layers = num_linear_layers - num_fp32_modules num_quantized_layers = 0 for name, module in model_quantized.named_modules(): if isinstance(module, torch.nn.Linear): if hasattr(model_quantized, "_keep_in_fp32_modules") and model_quantized._keep_in_fp32_modules: if any(fp32_name in name for fp32_name in model_quantized._keep_in_fp32_modules): continue self._verify_if_layer_quantized(name, module, config_kwargs) num_quantized_layers += 1 assert num_quantized_layers > 0, ( f"No quantized layers found in model (expected {expected_quantized_layers} linear layers, {num_fp32_modules} kept in FP32)" ) assert num_quantized_layers == expected_quantized_layers, ( f"Quantized layer count mismatch: expected {expected_quantized_layers}, got {num_quantized_layers} (total linear layers: {num_linear_layers}, FP32 modules: {num_fp32_modules})" ) def _test_keep_modules_in_fp32(self, config_kwargs): """ Test that modules listed in the model's `_keep_in_fp32_modules` stay in FP32 after quantization. Args: config_kwargs: Quantization config parameters """ fp32_modules = getattr(self.model_class, "_keep_in_fp32_modules", None) if not fp32_modules: pytest.skip(f"{self.model_class.__name__} does not declare _keep_in_fp32_modules") model = self._create_quantized_model(config_kwargs) model.to(torch_device) for name, module in model.named_modules(): if isinstance(module, torch.nn.Linear) and any(fp32_name in name for fp32_name in fp32_modules): assert module.weight.dtype == torch.float32, ( f"Module {name} should be FP32 but is {module.weight.dtype}" ) @torch.no_grad() def _test_quantization_modules_to_not_convert( self, config_kwargs, modules_to_not_convert, exclusion_key="modules_to_not_convert" ): """ Test that modules specified in modules_to_not_convert are not quantized. Args: config_kwargs: Base quantization config kwargs modules_to_not_convert: List of module names to exclude from quantization exclusion_key: Name of the config parameter carrying the exclusion list (BitsAndBytesConfig calls it `llm_int8_skip_modules`) """ # Create config with the exclusion list config_kwargs_with_exclusion = config_kwargs.copy() config_kwargs_with_exclusion[exclusion_key] = modules_to_not_convert model_with_exclusion = self._create_quantized_model(config_kwargs_with_exclusion) # Find a module that should NOT be quantized found_excluded = False for name, module in model_with_exclusion.named_modules(): if isinstance(module, torch.nn.Linear): # Check if this module is in the exclusion list if any(excluded in name for excluded in modules_to_not_convert): found_excluded = True # This module should NOT be quantized assert not self._is_module_quantized(module, config_kwargs), ( f"Module {name} should not be quantized but was found to be quantized" ) assert found_excluded, f"No linear layers found in excluded modules: {modules_to_not_convert}" # Find a module that SHOULD be quantized (not in exclusion list) found_quantized = False for name, module in model_with_exclusion.named_modules(): if isinstance(module, torch.nn.Linear): # Check if this module is NOT in the exclusion list if not any(excluded in name for excluded in modules_to_not_convert): if self._is_module_quantized(module, config_kwargs): found_quantized = True break assert found_quantized, "No quantized layers found outside of excluded modules" # Inference must work on the mixed model: excluded modules run in the compute dtype next to # quantized ones (excluded linears do strict-dtype matmuls). model_with_exclusion.to(torch_device) output = model_with_exclusion(**self.get_dummy_inputs(), return_dict=False)[0] assert output is not None, "Model output is None" assert not torch.isnan(output).any(), "Model output contains NaN" # Compare memory footprint with fully quantized model model_fully_quantized = self._create_quantized_model(config_kwargs) mem_with_exclusion = model_with_exclusion.get_memory_footprint() mem_fully_quantized = model_fully_quantized.get_memory_footprint() assert mem_with_exclusion > mem_fully_quantized, ( f"Model with exclusions should be larger. With exclusion: {mem_with_exclusion}, fully quantized: {mem_fully_quantized}" ) @torch.no_grad() def _test_quantization_device_map(self, config_kwargs): """ Test that quantized models work correctly with device_map="auto". Args: config_kwargs: Base quantization config kwargs """ model = self._create_quantized_model(config_kwargs, device_map="auto") 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 kind, named_tensors in (("parameter", model.named_parameters()), ("buffer", model.named_buffers())): for name, tensor in named_tensors: assert tensor.device.type != "meta", f"{kind} {name} was left on the meta device" if len(map_devices) == 1: assert tensor.device.type == next(iter(map_devices)), ( f"Expected device {next(iter(map_devices))} for {kind} {name}, got {tensor.device}" ) 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" def _test_quantization_cpu_device_map(self, config_kwargs): """ Test that quantized models are placed on the CPU with device_map="cpu". Args: config_kwargs: Base quantization config kwargs """ model_quantized = self._create_quantized_model(config_kwargs, device_map="cpu") assert hasattr(model_quantized, "hf_device_map"), "Model should have hf_device_map attribute" assert model_quantized.hf_device_map is not None, "hf_device_map should not be None" assert model_quantized.device == torch.device("cpu"), ( f"Model should be on CPU, but is on {model_quantized.device}" ) @torch.no_grad() def _test_dequantize(self, config_kwargs): """ Test that dequantize() converts quantized model back to standard linear layers. Args: config_kwargs: Quantization config parameters """ model = self._create_quantized_model(config_kwargs) model.to(torch_device) if not hasattr(model, "dequantize"): pytest.skip("Model does not have dequantize method") model.dequantize() for name, module in model.named_modules(): if isinstance(module, torch.nn.Linear): assert not self._is_module_quantized(module, config_kwargs), ( f"Module {name} is still quantized after dequantize()" ) inputs = self.get_dummy_inputs() output = model(**inputs, return_dict=False)[0] assert output is not None, "Model output is None after dequantization" assert not torch.isnan(output).any(), "Model output contains NaN after dequantization" def _test_quantization_training(self, config_kwargs): """ Test that quantized models can be used for training with LoRA-like adapters. This test: 1. Freezes all model parameters 2. Casts small parameters (e.g., layernorm) to fp32 for stability 3. Adds LoRA adapters to attention layers 4. Runs forward and backward passes 5. Verifies gradients are computed correctly Args: config_kwargs: Quantization config parameters """ model = self._create_quantized_model(config_kwargs) # Step 1: freeze all parameters for param in model.parameters(): param.requires_grad = False if param.ndim == 1: # cast small parameters (e.g. layernorm) to fp32 for stability param.data = param.data.to(torch.float32) # Step 2: add adapters to attention layers adapter_count = 0 for _, module in model.named_modules(): if "Attention" in repr(type(module)): if hasattr(module, "to_k"): module.to_k = LoRALayer(module.to_k, rank=4) adapter_count += 1 if hasattr(module, "to_q"): module.to_q = LoRALayer(module.to_q, rank=4) adapter_count += 1 if hasattr(module, "to_v"): module.to_v = LoRALayer(module.to_v, rank=4) adapter_count += 1 if adapter_count == 0: pytest.skip("No attention layers found in model for adapter training test") # Step 3: run forward and backward pass inputs = self.get_dummy_inputs() # Use bfloat16 on XPU and for bfloat16 models to avoid gradient underflow with quantized layers use_bf16 = torch_device == "xpu" or getattr(self, "torch_dtype", None) == torch.bfloat16 autocast_dtype = torch.bfloat16 if use_bf16 else torch.float16 with torch.amp.autocast(torch_device, dtype=autocast_dtype): out = model(**inputs, return_dict=False)[0] out.norm().backward() # Step 4: verify gradients are computed for module in model.modules(): if isinstance(module, LoRALayer): assert module.adapter[1].weight.grad is not None, "LoRA adapter gradient is None" assert module.adapter[1].weight.grad.norm().item() > 0, "LoRA adapter gradient norm is zero" @is_quantization @is_bitsandbytes @require_accelerator @require_bitsandbytes_version_greater("0.43.2") @require_accelerate class BitsAndBytesConfigMixin: """ Base mixin providing BitsAndBytes quantization config and model creation. Expected class attributes: - model_class: The model class to test - pretrained_model_name_or_path: Hub repository ID for the pretrained model - pretrained_model_kwargs: (Optional) Dict of kwargs to pass to from_pretrained """ BNB_CONFIGS = { "4bit_nf4": { "load_in_4bit": True, "bnb_4bit_quant_type": "nf4", "bnb_4bit_compute_dtype": torch.float16, }, "4bit_fp4": { "load_in_4bit": True, "bnb_4bit_quant_type": "fp4", "bnb_4bit_compute_dtype": torch.float16, }, "8bit": { "load_in_8bit": True, }, } BNB_EXPECTED_MEMORY_REDUCTIONS = { "4bit_nf4": 3.0, "4bit_fp4": 3.0, "8bit": 1.5, } sharded_serialization_config = BNB_CONFIGS["8bit"] def _create_quantized_model(self, config_kwargs, **extra_kwargs): config = BitsAndBytesConfig(**config_kwargs) kwargs = getattr(self, "pretrained_model_kwargs", {}).copy() kwargs["quantization_config"] = config kwargs.update(extra_kwargs) return self.model_class.from_pretrained(self.pretrained_model_name_or_path, **kwargs) def _verify_if_layer_quantized(self, name, module, config_kwargs): expected_weight_class = bnb.nn.Params4bit if config_kwargs.get("load_in_4bit") else bnb.nn.Int8Params assert module.weight.__class__ == expected_weight_class, ( f"Layer {name} has weight type {module.weight.__class__}, expected {expected_weight_class}" ) @is_bitsandbytes @require_accelerator @require_bitsandbytes_version_greater("0.43.2") @require_accelerate class BitsAndBytesTesterMixin(BitsAndBytesConfigMixin, QuantizationTesterMixin): """ Mixin class for testing BitsAndBytes quantization on models. Expected class attributes: - model_class: The model class to test - pretrained_model_name_or_path: Hub repository ID for the pretrained model - pretrained_model_kwargs: (Optional) Dict of kwargs to pass to from_pretrained (e.g., {"subfolder": "transformer"}) Expected methods to be implemented by subclasses: - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass Optional class attributes: - BNB_CONFIGS: Dict of config name -> BitsAndBytesConfig kwargs to test Pytest mark: bitsandbytes Use `pytest -m "not bitsandbytes"` to skip these tests """ @pytest.mark.parametrize( "config_name", list(BitsAndBytesConfigMixin.BNB_CONFIGS.keys()), ids=list(BitsAndBytesConfigMixin.BNB_CONFIGS.keys()), ) def test_bnb_quantization_num_parameters(self, config_name): self._test_quantization_num_parameters(BitsAndBytesConfigMixin.BNB_CONFIGS[config_name]) @pytest.mark.parametrize( "config_name", list(BitsAndBytesConfigMixin.BNB_CONFIGS.keys()), ids=list(BitsAndBytesConfigMixin.BNB_CONFIGS.keys()), ) def test_bnb_quantization_memory_footprint(self, config_name): expected = self.BNB_EXPECTED_MEMORY_REDUCTIONS.get(config_name, 1.2) self._test_quantization_memory_footprint( BitsAndBytesConfigMixin.BNB_CONFIGS[config_name], expected_memory_reduction=expected ) @pytest.mark.parametrize( "config_name", list(BitsAndBytesConfigMixin.BNB_CONFIGS.keys()), ids=list(BitsAndBytesConfigMixin.BNB_CONFIGS.keys()), ) def test_bnb_quantization_inference(self, config_name): self._test_quantization_inference(BitsAndBytesConfigMixin.BNB_CONFIGS[config_name]) @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", 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) @pytest.mark.parametrize( "config_name", list(BitsAndBytesConfigMixin.BNB_CONFIGS.keys()), ids=list(BitsAndBytesConfigMixin.BNB_CONFIGS.keys()), ) def test_bnb_quantized_layers(self, config_name): self._test_quantized_layers(BitsAndBytesConfigMixin.BNB_CONFIGS[config_name]) @pytest.mark.parametrize( "config_name", list(BitsAndBytesConfigMixin.BNB_CONFIGS.keys()), ids=list(BitsAndBytesConfigMixin.BNB_CONFIGS.keys()), ) def test_bnb_quantization_config_serialization(self, config_name): self._test_quantization_config_serialization(BitsAndBytesConfigMixin.BNB_CONFIGS[config_name]) def test_bnb_original_dtype(self): self._test_original_dtype(BitsAndBytesConfigMixin.BNB_CONFIGS["4bit_nf4"]) def test_bnb_keep_modules_in_fp32(self): self._test_keep_modules_in_fp32(BitsAndBytesConfigMixin.BNB_CONFIGS["4bit_nf4"]) @pytest.mark.parametrize("config_name", ["4bit_nf4", "8bit"], ids=["4bit_nf4", "8bit"]) def test_bnb_modules_to_not_convert(self, config_name): """Test module exclusion, which BitsAndBytesConfig exposes as `llm_int8_skip_modules` (despite the name, it also applies to 4-bit quantization).""" 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[config_name], modules_to_exclude, exclusion_key="llm_int8_skip_modules", ) 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.""" self._test_quantization_device_map(BitsAndBytesConfigMixin.BNB_CONFIGS[config_name]) def test_bnb_dequantize(self): """Test that dequantize() works correctly.""" self._test_dequantize(BitsAndBytesConfigMixin.BNB_CONFIGS["4bit_nf4"]) @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[config_name]) @pytest.mark.parametrize( "config_name", list(BitsAndBytesConfigMixin.BNB_CONFIGS.keys()), ids=list(BitsAndBytesConfigMixin.BNB_CONFIGS.keys()), ) def test_bnb_cpu_device_map(self, config_name): """Test that device_map='cpu' works correctly with quantization.""" self._test_quantization_cpu_device_map(BitsAndBytesConfigMixin.BNB_CONFIGS[config_name]) @is_quantization @is_torchao @require_accelerator @require_torchao_version_greater_or_equal("0.7.0") class TorchAoConfigMixin: """ Base mixin providing TorchAO quantization config and model creation. Expected class attributes: - model_class: The model class to test - pretrained_model_name_or_path: Hub repository ID for the pretrained model - pretrained_model_kwargs: (Optional) Dict of kwargs to pass to from_pretrained """ TORCHAO_QUANT_TYPES = { "int4wo": "Int4WeightOnlyConfig", "int8wo": "Int8WeightOnlyConfig", "int8dq": "Int8DynamicActivationInt8WeightConfig", } TORCHAO_EXPECTED_MEMORY_REDUCTIONS = { "int4wo": 1.8, "int8wo": 1.5, "int8dq": 1.5, } @staticmethod def _get_quant_config(config_name, modules_to_not_convert=None): # Quant types that need constructor arguments (e.g. `FqnToConfig`) are passed in already built. if not isinstance(config_name, str): return TorchAoConfig(config_name, modules_to_not_convert=modules_to_not_convert) config_cls = getattr(_torchao_quantization, config_name) config_kwargs = {"version": 2} # version=2 int4 defaults to the "plain" packing format, which routes through the # fbgemm/mslk Int4Tensor kernels. Pin the packing format to the tinygemm # (_convert_weight_to_int4pack) path on CUDA and plain_int32 on Intel XPU so the tests # don't require those extra kernels to be installed. if config_name == "Int4WeightOnlyConfig": if torch_device == "xpu": config_kwargs["int4_packing_format"] = "plain_int32" elif torch_device == "cuda": config_kwargs["int4_packing_format"] = "tile_packed_to_4d" return TorchAoConfig(config_cls(**config_kwargs), modules_to_not_convert=modules_to_not_convert) def _create_quantized_model(self, config_name, modules_to_not_convert=None, **extra_kwargs): config = self._get_quant_config(config_name, modules_to_not_convert=modules_to_not_convert) kwargs = getattr(self, "pretrained_model_kwargs", {}).copy() kwargs["quantization_config"] = config kwargs["device_map"] = str(torch_device) kwargs.update(extra_kwargs) return self.model_class.from_pretrained(self.pretrained_model_name_or_path, **kwargs) def _verify_if_layer_quantized(self, name, module, config_kwargs): from torchao.utils import TorchAOBaseTensor assert isinstance(module, torch.nn.Linear), f"Layer {name} is not Linear, got {type(module)}" assert isinstance(module.weight, TorchAOBaseTensor), ( f"Layer {name} weight is {type(module.weight)}, expected TorchAOBaseTensor" ) # int4wo requires CUDA or XPU ops (_convert_weight_to_int4pack) _int4wo_skip = pytest.mark.skipif( torch_device not in ["cuda", "xpu"], reason="int4wo quantization requires CUDA or XPU" ) # The CUDA int4 tinygemm path (Int4TilePackedTo4dTensor) does not implement aten.dequantize. _int4wo_dequantize_skip = pytest.mark.skip( reason="int4wo tinygemm packing (Int4TilePackedTo4dTensor) does not support dequantize" ) @is_torchao @require_accelerator @require_torchao_version_greater_or_equal("0.7.0") class TorchAoTesterMixin(TorchAoConfigMixin, QuantizationTesterMixin): """ Mixin class for testing TorchAO quantization on models. Expected class attributes: - model_class: The model class to test - pretrained_model_name_or_path: Hub repository ID for the pretrained model - pretrained_model_kwargs: (Optional) Dict of kwargs to pass to from_pretrained (e.g., {"subfolder": "transformer"}) Expected methods to be implemented by subclasses: - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass Optional class attributes: - TORCHAO_QUANT_TYPES: Dict of quantization type strings to test Pytest mark: torchao Use `pytest -m "not torchao"` to skip these tests """ @pytest.mark.parametrize( "quant_type", [ pytest.param("int4wo", marks=_int4wo_skip), "int8wo", "int8dq", ], ids=["int4wo", "int8wo", "int8dq"], ) def test_torchao_quantization_num_parameters(self, quant_type): self._test_quantization_num_parameters(TorchAoConfigMixin.TORCHAO_QUANT_TYPES[quant_type]) @pytest.mark.parametrize( "quant_type", [ pytest.param("int4wo", marks=_int4wo_skip), "int8wo", "int8dq", ], ids=["int4wo", "int8wo", "int8dq"], ) def test_torchao_quantization_memory_footprint(self, quant_type): expected = TorchAoConfigMixin.TORCHAO_EXPECTED_MEMORY_REDUCTIONS.get(quant_type, 1.2) self._test_quantization_memory_footprint( TorchAoConfigMixin.TORCHAO_QUANT_TYPES[quant_type], expected_memory_reduction=expected ) @pytest.mark.parametrize( "quant_type", [ pytest.param("int4wo", marks=[_int4wo_skip, _int4wo_dequantize_skip]), "int8wo", "int8dq", ], ids=["int4wo", "int8wo", "int8dq"], ) def test_torchao_quantization_inference(self, quant_type): self._test_quantization_inference(TorchAoConfigMixin.TORCHAO_QUANT_TYPES[quant_type]) @pytest.mark.parametrize("quant_type", ["int8wo"], ids=["int8wo"]) def test_torchao_quantized_layers(self, quant_type): self._test_quantized_layers(TorchAoConfigMixin.TORCHAO_QUANT_TYPES[quant_type]) @pytest.mark.parametrize("quant_type", ["int8wo"], ids=["int8wo"]) def test_torchao_quantization_lora_inference(self, quant_type): self._test_quantization_lora_inference(TorchAoConfigMixin.TORCHAO_QUANT_TYPES[quant_type]) @pytest.mark.parametrize("quant_type", ["int8wo"], ids=["int8wo"]) @require_torchao_version_greater_or_equal("0.16.0") def test_torchao_quantization_serialization(self, quant_type, tmp_path): self._test_quantization_serialization(TorchAoConfigMixin.TORCHAO_QUANT_TYPES[quant_type], tmp_path) @pytest.mark.parametrize("quant_type", ["int8dq"], ids=["int8dq"]) @require_torchao_version_greater_or_equal("0.16.0") def test_torchao_quantization_sharded_serialization(self, quant_type, tmp_path): self._test_quantization_serialization( TorchAoConfigMixin.TORCHAO_QUANT_TYPES[quant_type], tmp_path, max_shard_size="16KB" ) def test_torchao_modules_to_not_convert(self): """Test that modules_to_not_convert parameter 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") # TorchAoConfig takes modules_to_not_convert directly (not inside the quant_type config), # so this can't reuse the dict-based QuantizationTesterMixin helper. model = self._create_quantized_model( TorchAoConfigMixin.TORCHAO_QUANT_TYPES["int8wo"], modules_to_not_convert=modules_to_exclude ) found_excluded = False for name, module in model.named_modules(): if isinstance(module, torch.nn.Linear) and any(excluded in name for excluded in modules_to_exclude): found_excluded = True assert not self._is_module_quantized(module), ( f"Module {name} should not be quantized but was found to be quantized" ) assert found_excluded, f"No linear layers found in excluded modules: {modules_to_exclude}" # Inference must work on the mixed model (see _test_quantization_modules_to_not_convert). with torch.no_grad(): output = model(**self.get_dummy_inputs(), return_dict=False)[0] assert output is not None, "Model output is None" assert not torch.isnan(output).any(), "Model output contains NaN" 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. Root-level parameters and buffers (e.g. Wan's `scale_shift_table`) need their own # device-map entries since they belong to no child module. 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) for name, _ in list(model.named_parameters(recurse=False)) + list(model.named_buffers(recurse=False)): 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 (some of) the weights on the accelerator # are. Not every accelerator module is necessarily quantized: the offload exclusion matches # module names by substring, so an offloaded `blocks` also excludes e.g. Wan's `vace_blocks`. found_quantized = False 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" ) elif isinstance(module.weight, TorchAOBaseTensor): found_quantized = True assert found_quantized, "No quantized layers found outside the offloaded modules" 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" def test_torchao_dequantize(self): """Test that dequantize() works correctly.""" self._test_dequantize(TorchAoConfigMixin.TORCHAO_QUANT_TYPES["int8wo"]) def test_torchao_training(self): """Test that quantized models can be used for training with adapters.""" self._test_quantization_training(TorchAoConfigMixin.TORCHAO_QUANT_TYPES["int8wo"]) def test_torchao_keep_modules_in_fp32(self): self._test_keep_modules_in_fp32(TorchAoConfigMixin.TORCHAO_QUANT_TYPES["int8wo"]) def test_torchao_fqn_to_config(self): """An `FqnToConfig` quant type must reach `quantize_` instead of tripping its `filter_fn` guard. See https://github.com/huggingface/diffusers/issues/14667. """ quant_type = _torchao_quantization.FqnToConfig( {"_default": _torchao_quantization.Int8WeightOnlyConfig(version=2)} ) self._test_quantized_layers(quant_type) def test_torchao_fqn_to_config_targets_named_layers(self): """Selective `FqnToConfig` targeting must land on the same layers as a plain `quantize_` call. The quantizer resolves fqns one `nn.Linear` at a time, so `quantization_config=` is checked against torchao's own whole-model pass. See https://github.com/huggingface/diffusers/issues/14667. """ keep_in_fp32 = getattr(self.model_class, "_keep_in_fp32_modules", None) or [] reference = self._load_unquantized_model().to(torch_device) linear_fqns = [ name for name, module in reference.named_modules() if isinstance(module, torch.nn.Linear) and not any(fp32_name in name for fp32_name in keep_in_fp32) ] if len(linear_fqns) < 2: pytest.skip("Model does not have enough linear layers to test selective fqn quantization") # One layer selected by its exact fqn, a set of them by a wildcard regex over the model's own fqns # (the shape from the issue https://github.com/huggingface/diffusers/issues/14667), everything else # left alone by `_default: None`. On Flux2, this gives # exact_fqn="time_guidance_embed.timestep_embedder.linear_1" and # pattern=".*proj_out" (matching "proj_out"). exact_fqn = linear_fqns[0] pattern = f".*{re.escape(linear_fqns[-1].split('.')[-1])}" fqn_to_config = { exact_fqn: _torchao_quantization.Int8WeightOnlyConfig(version=2), f"re:{pattern}": _torchao_quantization.Int8WeightOnlyConfig(version=2), "_default": None, } # `quantize_` on the whole model is the reference: `quantization_config=` must place the same tensors. _torchao_quantization.quantize_(reference, _torchao_quantization.FqnToConfig(fqn_to_config), filter_fn=None) expected = {name: type(reference.get_submodule(name).weight) for name in linear_fqns} del reference assert set(expected.values()) != {torch.nn.Parameter}, "`quantize_` did not quantize any of the named layers" model = self._create_quantized_model(_torchao_quantization.FqnToConfig(fqn_to_config)) quantized = {name: type(model.get_submodule(name).weight) for name in linear_fqns} mismatched = { name: (quantized[name], expected[name]) for name in linear_fqns if quantized[name] != expected[name] } assert not mismatched, f"Layers quantized through `quantization_config=` differ from `quantize_`: {mismatched}" @is_quantization @is_gguf @require_accelerate @require_accelerator @require_gguf_version_greater_or_equal("0.10.0") class GGUFConfigMixin: """ Base mixin providing GGUF quantization config and model creation. Expected from config mixin: - model_class: The model class to test Required properties (must be implemented by subclasses): - gguf_filename: URL or path to the GGUF file """ @property def gguf_filename(self): """URL or path to the GGUF file. Must be implemented by subclasses.""" raise NotImplementedError("Subclasses must implement the `gguf_filename` property.") def _create_quantized_model(self, config_kwargs=None, **extra_kwargs): if config_kwargs is None: config_kwargs = {"compute_dtype": torch.bfloat16} config = GGUFQuantizationConfig(**config_kwargs) kwargs = { "quantization_config": config, "torch_dtype": config_kwargs.get("compute_dtype", torch.bfloat16), "device_map": str(torch_device), } kwargs.update(extra_kwargs) return self.model_class.from_single_file(self.gguf_filename, **kwargs) def _verify_if_layer_quantized(self, name, module, config_kwargs=None): from diffusers.quantizers.gguf.utils import GGUFParameter assert isinstance(module.weight, GGUFParameter), f"{name} weight is not GGUFParameter" assert hasattr(module.weight, "quant_type"), f"{name} weight missing quant_type" assert module.weight.dtype == torch.uint8, f"{name} weight dtype should be uint8" @is_gguf @require_accelerate @require_accelerator @require_gguf_version_greater_or_equal("0.10.0") class GGUFTesterMixin(GGUFConfigMixin, QuantizationTesterMixin): """ Mixin class for testing GGUF quantization on models. Expected from config mixin: - model_class: The model class to test Required properties (must be implemented by subclasses): - gguf_filename: URL or path to the GGUF file Expected methods from config mixin: - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass Pytest mark: gguf Use `pytest -m "not gguf"` to skip these tests """ def test_gguf_quantization_inference(self): self._test_quantization_inference({"compute_dtype": torch.bfloat16}) 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") _keep_in_fp32_modules = self.model_class._keep_in_fp32_modules self.model_class._keep_in_fp32_modules = ["proj_out"] try: self._test_keep_modules_in_fp32({"compute_dtype": torch.bfloat16}) finally: self.model_class._keep_in_fp32_modules = _keep_in_fp32_modules def test_gguf_quantization_dtype_assignment(self): self._test_quantization_dtype_assignment({"compute_dtype": torch.bfloat16}) def test_gguf_quantization_lora_inference(self): self._test_quantization_lora_inference({"compute_dtype": torch.bfloat16}) def test_gguf_dequantize(self): """Test that dequantize() works correctly.""" self._test_dequantize({"compute_dtype": torch.bfloat16}) @is_quantization @is_modelopt @require_accelerator @require_accelerate @require_modelopt_version_greater_or_equal("0.33.1") class ModelOptConfigMixin: """ Base mixin providing NVIDIA ModelOpt quantization config and model creation. Expected class attributes: - model_class: The model class to test - pretrained_model_name_or_path: Hub repository ID for the pretrained model - pretrained_model_kwargs: (Optional) Dict of kwargs to pass to from_pretrained """ MODELOPT_CONFIGS = { "fp8": {"quant_type": "FP8"}, "int8": {"quant_type": "INT8"}, "int4": {"quant_type": "INT4"}, } MODELOPT_EXPECTED_MEMORY_REDUCTIONS = { "fp8": 1.5, "int8": 1.5, "int4": 3.0, } def _create_quantized_model(self, config_kwargs, **extra_kwargs): config = NVIDIAModelOptConfig(**config_kwargs) kwargs = getattr(self, "pretrained_model_kwargs", {}).copy() kwargs["quantization_config"] = config kwargs["device_map"] = str(torch_device) kwargs.update(extra_kwargs) return self.model_class.from_pretrained(self.pretrained_model_name_or_path, **kwargs) def _verify_if_layer_quantized(self, name, module, config_kwargs): assert mtq.utils.is_quantized(module), f"Layer {name} does not have weight_quantizer attribute (not quantized)" @is_modelopt @require_accelerator @require_accelerate @require_modelopt_version_greater_or_equal("0.33.1") class ModelOptTesterMixin(ModelOptConfigMixin, QuantizationTesterMixin): """ Mixin class for testing NVIDIA ModelOpt quantization on models. Expected class attributes: - model_class: The model class to test - pretrained_model_name_or_path: Hub repository ID for the pretrained model - pretrained_model_kwargs: (Optional) Dict of kwargs to pass to from_pretrained (e.g., {"subfolder": "transformer"}) Expected methods to be implemented by subclasses: - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass Optional class attributes: - MODELOPT_CONFIGS: Dict of config name -> NVIDIAModelOptConfig kwargs to test Pytest mark: modelopt Use `pytest -m "not modelopt"` to skip these tests """ @pytest.mark.parametrize("config_name", ["fp8"], ids=["fp8"]) def test_modelopt_quantization_num_parameters(self, config_name): self._test_quantization_num_parameters(ModelOptConfigMixin.MODELOPT_CONFIGS[config_name]) @pytest.mark.parametrize( "config_name", list(ModelOptConfigMixin.MODELOPT_CONFIGS.keys()), ids=list(ModelOptConfigMixin.MODELOPT_CONFIGS.keys()), ) def test_modelopt_quantization_memory_footprint(self, config_name): expected = ModelOptConfigMixin.MODELOPT_EXPECTED_MEMORY_REDUCTIONS.get(config_name, 1.2) self._test_quantization_memory_footprint( ModelOptConfigMixin.MODELOPT_CONFIGS[config_name], expected_memory_reduction=expected ) @pytest.mark.parametrize( "config_name", list(ModelOptConfigMixin.MODELOPT_CONFIGS.keys()), ids=list(ModelOptConfigMixin.MODELOPT_CONFIGS.keys()), ) def test_modelopt_quantization_inference(self, config_name): self._test_quantization_inference(ModelOptConfigMixin.MODELOPT_CONFIGS[config_name]) @pytest.mark.parametrize("config_name", ["fp8"], ids=["fp8"]) def test_modelopt_quantization_dtype_assignment(self, config_name): self._test_quantization_dtype_assignment(ModelOptConfigMixin.MODELOPT_CONFIGS[config_name]) @pytest.mark.parametrize("config_name", ["fp8"], ids=["fp8"]) def test_modelopt_quantization_lora_inference(self, config_name): self._test_quantization_lora_inference(ModelOptConfigMixin.MODELOPT_CONFIGS[config_name]) @pytest.mark.parametrize("config_name", ["fp8"], ids=["fp8"]) def test_modelopt_quantization_serialization(self, config_name, tmp_path): self._test_quantization_serialization(ModelOptConfigMixin.MODELOPT_CONFIGS[config_name], tmp_path) @pytest.mark.parametrize("config_name", ["fp8"], ids=["fp8"]) def test_modelopt_quantized_layers(self, config_name): self._test_quantized_layers(ModelOptConfigMixin.MODELOPT_CONFIGS[config_name]) def test_modelopt_modules_to_not_convert(self): """Test that modules_to_not_convert parameter 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(ModelOptConfigMixin.MODELOPT_CONFIGS["fp8"], modules_to_exclude) def test_modelopt_device_map(self): """Test that device_map='auto' works correctly with quantization.""" self._test_quantization_device_map(ModelOptConfigMixin.MODELOPT_CONFIGS["fp8"]) def test_modelopt_dequantize(self): """Test that dequantize() works correctly.""" self._test_dequantize(ModelOptConfigMixin.MODELOPT_CONFIGS["fp8"]) @is_quantization @is_sdnq @require_sdnq @require_accelerate @require_accelerator class SDNQConfigMixin: """ Base mixin providing SDNQ quantization config and model creation. Expected class attributes: - model_class: The model class to test - pretrained_model_name_or_path: Hub repository ID for the pretrained model - pretrained_model_kwargs: (Optional) Dict of kwargs to pass to from_pretrained """ SDNQ_CONFIGS = { "int8": {"weights_dtype": "int8"}, "uint4_svd": {"weights_dtype": "uint4", "use_svd": True}, } SDNQ_EXPECTED_MEMORY_REDUCTIONS = { "int8": 1.5, "uint4_svd": 1.5, } def _create_quantized_model(self, config_kwargs, **extra_kwargs): config = SDNQConfig(**config_kwargs) kwargs = getattr(self, "pretrained_model_kwargs", {}).copy() kwargs["quantization_config"] = config # SDNQ leaves some modules (embedders, norms) unquantized, so the compute dtype must be set at load # time to keep them consistent with the quantized layers. kwargs.setdefault("dtype", getattr(self, "torch_dtype", torch.bfloat16)) kwargs.update(extra_kwargs) return self.model_class.from_pretrained(self.pretrained_model_name_or_path, **kwargs) def _verify_if_layer_quantized(self, name, module, config_kwargs): assert hasattr(module, "sdnq_dequantizer"), f"Layer {name} is not an SDNQ quantized layer" @is_sdnq @require_sdnq @require_accelerate @require_accelerator class SDNQTesterMixin(SDNQConfigMixin, QuantizationTesterMixin): """ Mixin class for testing SDNQ quantization on models. Expected class attributes: - model_class: The model class to test - pretrained_model_name_or_path: Hub repository ID for the pretrained model - pretrained_model_kwargs: (Optional) Dict of kwargs to pass to from_pretrained (e.g., {"subfolder": "transformer"}) Expected methods to be implemented by subclasses: - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass Optional class attributes: - SDNQ_CONFIGS: Dict of config name -> SDNQConfig kwargs to test Pytest mark: sdnq Use `pytest -m "not sdnq"` to skip these tests """ @pytest.mark.parametrize( "config_name", list(SDNQConfigMixin.SDNQ_CONFIGS.keys()), ids=list(SDNQConfigMixin.SDNQ_CONFIGS.keys()), ) def test_sdnq_quantization_inference(self, config_name): self._test_quantization_inference(SDNQConfigMixin.SDNQ_CONFIGS[config_name]) @pytest.mark.parametrize( "config_name", list(SDNQConfigMixin.SDNQ_CONFIGS.keys()), ids=list(SDNQConfigMixin.SDNQ_CONFIGS.keys()), ) def test_sdnq_quantization_memory_footprint(self, config_name): expected = SDNQConfigMixin.SDNQ_EXPECTED_MEMORY_REDUCTIONS.get(config_name, 1.2) self._test_quantization_memory_footprint( SDNQConfigMixin.SDNQ_CONFIGS[config_name], expected_memory_reduction=expected ) @torch.no_grad() def test_sdnq_quantization_serialization(self, tmp_path): # Unlike the base helper, reload onto the accelerator: SDNQ reloads to CPU, while dummy inputs live on device. model = self._create_quantized_model(SDNQConfigMixin.SDNQ_CONFIGS["int8"]) model.save_pretrained(str(tmp_path), safe_serialization=True) # No quantization_config passed, so it must be picked up from the saved config.json. model_loaded = self.model_class.from_pretrained(str(tmp_path)).to(torch_device) inputs = self.get_dummy_inputs() output = model_loaded(**inputs, return_dict=False)[0] assert not torch.isnan(output).any(), "Loaded model output contains NaN" @pytest.mark.parametrize("config_name", ["int8"], ids=["int8"]) def test_sdnq_quantization_dtype_assignment(self, config_name): self._test_quantization_dtype_assignment(SDNQConfigMixin.SDNQ_CONFIGS[config_name]) def test_sdnq_modules_to_not_convert(self): """Test that modules_to_not_convert parameter 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(SDNQConfigMixin.SDNQ_CONFIGS["int8"], modules_to_exclude) def test_sdnq_dequantize(self): """Test that dequantize() works correctly.""" self._test_dequantize(SDNQConfigMixin.SDNQ_CONFIGS["int8"]) @is_quantization @is_torch_compile class QuantizationCompileTesterMixin: """ Base mixin class providing common test implementations for torch.compile with quantized models. Backend-specific compile mixins should: 1. Inherit from their respective config mixin (e.g., BitsAndBytesConfigMixin) 2. Inherit from this mixin 3. Define the config to use for compile tests Expected class attributes: - model_class: The model class to test - pretrained_model_name_or_path: Hub repository ID for the pretrained model - pretrained_model_kwargs: (Optional) Dict of kwargs to pass to from_pretrained Expected methods in test classes: - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass """ def setup_method(self): gc.collect() backend_empty_cache(torch_device) torch.compiler.reset() def teardown_method(self): gc.collect() backend_empty_cache(torch_device) torch.compiler.reset() @torch.no_grad() def _test_torch_compile(self, config_kwargs, fullgraph=True, error_on_recompile=True): """ Test that torch.compile works correctly with a quantized model. Args: config_kwargs: Quantization config parameters fullgraph: Whether the compiled model is required to compile without graph breaks error_on_recompile: Whether a recompilation during the forward pass should fail the test """ model = self._create_quantized_model(config_kwargs) model.to(torch_device) model.eval() model.compile(fullgraph=fullgraph) with torch._dynamo.config.patch(error_on_recompile=error_on_recompile): 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" @torch.no_grad() def _test_torch_compile_with_group_offload(self, config_kwargs, use_stream=False): """ Test that torch.compile works correctly with a quantized model and group offloading. Args: config_kwargs: Quantization config parameters use_stream: Whether to use CUDA streams for offloading """ torch._dynamo.config.cache_size_limit = 1000 model = self._create_quantized_model(config_kwargs) model.eval() if not hasattr(model, "enable_group_offload"): pytest.skip("Model does not support group offloading") group_offload_kwargs = { "onload_device": torch.device(torch_device), "offload_device": torch.device("cpu"), "offload_type": "leaf_level", "use_stream": use_stream, } model.enable_group_offload(**group_offload_kwargs) model.compile() 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" @is_bitsandbytes @require_accelerator @require_bitsandbytes_version_greater("0.43.2") @require_accelerate class BitsAndBytesCompileTesterMixin(BitsAndBytesConfigMixin, QuantizationCompileTesterMixin): """ Mixin class for testing torch.compile with BitsAndBytes quantized models. Expected class attributes: - model_class: The model class to test - pretrained_model_name_or_path: Hub repository ID for the pretrained model - pretrained_model_kwargs: (Optional) Dict of kwargs to pass to from_pretrained Expected methods to be implemented by subclasses: - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass Pytest mark: bitsandbytes Use `pytest -m "not bitsandbytes"` to skip these tests """ @pytest.mark.parametrize("config_name", ["4bit_nf4"], ids=["4bit_nf4"]) def test_bnb_torch_compile(self, config_name): self._test_torch_compile(BitsAndBytesConfigMixin.BNB_CONFIGS[config_name]) @pytest.mark.parametrize("config_name", ["4bit_nf4"], ids=["4bit_nf4"]) def test_bnb_torch_compile_with_group_offload(self, config_name): self._test_torch_compile_with_group_offload(BitsAndBytesConfigMixin.BNB_CONFIGS[config_name]) @is_torchao @require_accelerator @require_torchao_version_greater_or_equal("0.7.0") class TorchAoCompileTesterMixin(TorchAoConfigMixin, QuantizationCompileTesterMixin): """ Mixin class for testing torch.compile with TorchAO quantized models. Expected class attributes: - model_class: The model class to test - pretrained_model_name_or_path: Hub repository ID for the pretrained model - pretrained_model_kwargs: (Optional) Dict of kwargs to pass to from_pretrained Expected methods to be implemented by subclasses: - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass Pytest mark: torchao Use `pytest -m "not torchao"` to skip these tests """ @pytest.mark.parametrize("quant_type", ["int8wo"], ids=["int8wo"]) def test_torchao_torch_compile(self, quant_type): self._test_torch_compile(TorchAoConfigMixin.TORCHAO_QUANT_TYPES[quant_type]) @pytest.mark.parametrize("quant_type", ["int8wo"], ids=["int8wo"]) def test_torchao_torch_compile_with_group_offload(self, quant_type): self._test_torch_compile_with_group_offload(TorchAoConfigMixin.TORCHAO_QUANT_TYPES[quant_type]) @is_gguf @require_accelerate @require_accelerator @require_gguf_version_greater_or_equal("0.10.0") class GGUFCompileTesterMixin(GGUFConfigMixin, QuantizationCompileTesterMixin): """ Mixin class for testing torch.compile with GGUF quantized models. Expected from config mixin: - model_class: The model class to test Required properties (must be implemented by subclasses): - gguf_filename: URL or path to the GGUF file Expected methods from config mixin: - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass Pytest mark: gguf Use `pytest -m "not gguf"` to skip these tests """ def test_gguf_torch_compile(self): self._test_torch_compile({"compute_dtype": torch.bfloat16}) def test_gguf_torch_compile_with_group_offload(self): self._test_torch_compile_with_group_offload({"compute_dtype": torch.bfloat16}) @pytest.mark.skipif(not is_kernels_available(), reason="`kernels` is not available.") @require_accelerate @require_accelerator class NunchakuLiteConfigMixin: """ Base mixin providing Nunchaku Lite quantization config and model creation. Expected class attributes: - model_class: The model class to test - quantized_model_name_or_path: Hub repository ID or local path for the quantized model - pretrained_model_kwargs: (Optional) Dict of kwargs to pass to from_pretrained """ config_dict = None def _create_quantized_model(self, config_kwargs=None, **extra_kwargs): kwargs = getattr(self, "pretrained_model_kwargs", {}).copy() if config_kwargs is not None: kwargs["quantization_config"] = NunchakuLiteQuantizationConfig(**config_kwargs) kwargs.update(extra_kwargs) return self.model_class.from_pretrained(self.quantized_model_name_or_path, **kwargs) def _verify_if_layer_quantized(self, name, module, config_kwargs): from diffusers.quantizers.nunchaku.utils import AWQW4A16Linear, SVDQW4A4Linear assert isinstance(module, (SVDQW4A4Linear, AWQW4A16Linear)), ( f"Layer {name} is not a Nunchaku Lite layer, got {type(module)}" ) @pytest.mark.skipif(not is_kernels_available(), reason="`kernels` is not available.") @require_accelerate @require_accelerator class NunchakuLiteTesterMixin(NunchakuLiteConfigMixin, QuantizationTesterMixin): """ Mixin class for testing Nunchaku Lite quantization on models. Expected class attributes: - model_class: The model class to test - quantized_model_name_or_path: Hub repository ID or local path for the quantized model - pretrained_model_kwargs: (Optional) Dict of kwargs to pass to from_pretrained Expected methods to be implemented by subclasses: - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass """ def test_nunchaku_lite_quantization_inference(self): self._test_quantization_inference(self.config_dict) def _is_module_quantized(self, module, config_kwargs=None): from diffusers.quantizers.nunchaku.utils import AWQW4A16Linear, SVDQW4A4Linear return isinstance(module, (SVDQW4A4Linear, AWQW4A16Linear)) def _test_quantized_layers(self, config_kwargs): model = self._create_quantized_model(config_kwargs) num_quantized_layers = 0 for name, module in model.named_modules(): if self._is_module_quantized(module): self._verify_if_layer_quantized(name, module, config_kwargs) num_quantized_layers += 1 expected_quantized_layers = sum( len(config_kwargs.get(section, {}).get("targets", [])) for section in ("svdq_w4a4", "awq_w4a16") ) assert num_quantized_layers > 0, ( f"No quantized layers found in model (expected {expected_quantized_layers} Nunchaku Lite layers)" ) assert num_quantized_layers == expected_quantized_layers, ( f"Quantized layer count mismatch: expected {expected_quantized_layers}, got {num_quantized_layers} " f"(configured Nunchaku Lite targets: {expected_quantized_layers})" ) def test_nunchaku_lite_quantized_layers(self): self._test_quantized_layers(self.config_dict) @pytest.mark.skipif(not is_kernels_available(), reason="`kernels` is not available.") @require_accelerate @require_accelerator class NunchakuLiteCompileTesterMixin(NunchakuLiteConfigMixin, QuantizationCompileTesterMixin): """ Mixin class for testing torch.compile with Nunchaku Lite quantized models. Expected class attributes: - model_class: The model class to test - quantized_model_name_or_path: Hub repository ID or local path for the quantized model - pretrained_model_kwargs: (Optional) Dict of kwargs to pass to from_pretrained Expected methods to be implemented by subclasses: - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass """ def test_nunchaku_lite_torch_compile(self): self._test_torch_compile(self.config_dict) def test_nunchaku_lite_torch_compile_with_group_offload(self): self._test_torch_compile_with_group_offload(self.config_dict) @is_modelopt @require_accelerator @require_accelerate @require_modelopt_version_greater_or_equal("0.33.1") class ModelOptCompileTesterMixin(ModelOptConfigMixin, QuantizationCompileTesterMixin): """ Mixin class for testing torch.compile with NVIDIA ModelOpt quantized models. Expected class attributes: - model_class: The model class to test - pretrained_model_name_or_path: Hub repository ID for the pretrained model - pretrained_model_kwargs: (Optional) Dict of kwargs to pass to from_pretrained Expected methods to be implemented by subclasses: - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass Pytest mark: modelopt Use `pytest -m "not modelopt"` to skip these tests """ @pytest.mark.parametrize("config_name", ["fp8"], ids=["fp8"]) def test_modelopt_torch_compile(self, config_name): self._test_torch_compile(ModelOptConfigMixin.MODELOPT_CONFIGS[config_name]) @pytest.mark.parametrize("config_name", ["fp8"], ids=["fp8"]) def test_modelopt_torch_compile_with_group_offload(self, config_name): self._test_torch_compile_with_group_offload(ModelOptConfigMixin.MODELOPT_CONFIGS[config_name]) @is_sdnq @require_sdnq @require_accelerate @require_accelerator class SDNQCompileTesterMixin(SDNQConfigMixin, QuantizationCompileTesterMixin): """ Mixin class for testing torch.compile with SDNQ quantized models. Expected class attributes: - model_class: The model class to test - pretrained_model_name_or_path: Hub repository ID for the pretrained model - pretrained_model_kwargs: (Optional) Dict of kwargs to pass to from_pretrained Expected methods to be implemented by subclasses: - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass Pytest mark: sdnq Use `pytest -m "not sdnq"` to skip these tests """ @pytest.mark.parametrize("config_name", ["int8"], ids=["int8"]) def test_sdnq_torch_compile(self, config_name): self._test_torch_compile(SDNQConfigMixin.SDNQ_CONFIGS[config_name]) @pytest.mark.parametrize("config_name", ["int8"], ids=["int8"]) def test_sdnq_torch_compile_with_group_offload(self, config_name): self._test_torch_compile_with_group_offload(SDNQConfigMixin.SDNQ_CONFIGS[config_name]) @is_quantization @is_autoround @require_accelerator @require_accelerate @require_auto_round_version_greater_or_equal("0.13.0") class AutoRoundConfigMixin: """ Base mixin providing AutoRound quantization config and model creation. AutoRound is a weight-only quantization method (W4A16). It supports multiple inference When `backend="auto"`, AutoRound selects the best available backend automatically. Expected class attributes: - model_class: The model class to test - pretrained_model_name_or_path: Hub repository ID for the pretrained model - quantized_model_name_or_path: Hub repository ID for the quantized model - pretrained_model_kwargs: (Optional) Dict of kwargs to pass to from_pretrained """ config_dict = {"backend": "auto"} def _create_quantized_model(self, config_kwargs, **extra_kwargs): config = AutoRoundConfig(**config_kwargs) kwargs = getattr(self, "pretrained_model_kwargs", {}).copy() kwargs["quantization_config"] = config kwargs["torch_dtype"] = torch.bfloat16 if "device_map" not in kwargs: kwargs["device_map"] = torch_device kwargs.update(extra_kwargs) return self.model_class.from_pretrained(self.quantized_model_name_or_path, **kwargs) def _verify_if_layer_quantized(self, name, module, config_kwargs): # AutoRound replaces linear layers with quantized linear layers assert isinstance(module, torch.nn.Linear), f"Layer {name} is not Linear, got {type(module)}" @is_autoround @require_accelerator @require_accelerate @require_auto_round_version_greater_or_equal("0.13.0") class AutoRoundTesterMixin(AutoRoundConfigMixin, QuantizationTesterMixin): """ Mixin class for testing AutoRound quantization on models. Expected class attributes: - model_class: The model class to test - pretrained_model_name_or_path: Hub repository ID for the pretrained model - quantized_model_name_or_path: Hub repository ID for the quantized model - pretrained_model_kwargs: (Optional) Dict of kwargs to pass to from_pretrained (e.g., {"subfolder": "transformer"}) Expected methods to be implemented by subclasses: - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass Optional class attributes: - AUTOROUND_CONFIGS: Dict of config name -> AutoRoundConfig kwargs to test Pytest mark: autoround Use `pytest -m "not autoround"` to skip these tests """ config_dict = {"backend": "auto"} def test_autoround_quantization_memory_footprint(self): expected = 1.5 # AutoRound is a W4A16 method, so we expect around 1.5x memory reduction self._test_quantization_memory_footprint(self.config_dict, expected_memory_reduction=expected) def test_autoround_quantization_inference(self): self._test_quantization_inference(self.config_dict) def test_autoround_device_map(self): """Test that device_map='auto' works correctly with quantization.""" self._test_quantization_device_map(self.config_dict) @is_autoround @require_accelerator @require_accelerate @require_auto_round_version_greater_or_equal("0.13.0") class AutoRoundCompileTesterMixin(AutoRoundConfigMixin, QuantizationCompileTesterMixin): """ Mixin class for testing `torch.compile` with AutoRound-quantized models. This mixin provides tests that verify `torch.compile` works correctly with models quantized using AutoRound. Subclasses are expected to inherit from `AutoRoundConfigMixin` (which defines `config_dict`) and to provide the following class attributes: `model_class`, `pretrained_model_name_or_path`, and `quantized_model_name_or_path`. The mixin uses `config_dict` (defaults to {"backend": "auto"}) as the quantization configuration passed into `_create_quantized_model` when invoking the compile-related tests. Provided tests: - `test_autoround_torch_compile`: Ensures `torch.compile` runs and produces valid, non-NaN outputs for an AutoRound-quantized model. - `test_autoround_torch_compile_with_group_offload`: Ensures `torch.compile` works together with group offloading when supported by the quantized model implementation. """ def test_autoround_torch_compile(self): self._test_torch_compile(self.config_dict, fullgraph=False, error_on_recompile=False) def test_autoround_torch_compile_with_group_offload(self): self._test_torch_compile_with_group_offload(self.config_dict) class BFLConfigMixin: """ Base mixin providing BFL (`fp8r` / `nvfp4`) quantization config and model creation. BFL checkpoints are prequantized single files, one per scheme, and the tests never download their weights. The safetensors header of each checkpoint says which linears the exporter quantized; the tests quantize those same linears of the small test model with the library's own `quantize_fp8r` / `quantize_nvfp4` and load the result through [`~FromOriginalModelMixin.from_single_file`], exactly as the real checkpoint is loaded. Expected class attributes: - model_class: The model class to test Expected methods to be implemented by subclasses: - get_init_dict(): Returns the init kwargs of the small test model - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass Required properties (must be implemented by subclasses): - bfl_checkpoints: Dict mapping a scheme (`"fp8r"`, `"nvfp4"`) to the URL or path of its single-file checkpoint """ BFL_CONFIGS = { "fp8r": {"scheme": "fp8r", "compute_dtype": torch.bfloat16}, "nvfp4": {"scheme": "nvfp4", "compute_dtype": torch.bfloat16}, } @property def bfl_checkpoints(self) -> dict[str, str]: """Scheme -> single-file checkpoint. Must be implemented by subclasses.""" raise NotImplementedError("Subclasses must implement the `bfl_checkpoints` property.") @staticmethod def _bfl_linear_cls(scheme): from diffusers.quantizers.bfl.utils import BFLFP8RLinear, BFLNVFP4Linear return {"fp8r": BFLFP8RLinear, "nvfp4": BFLNVFP4Linear}[scheme] def _bfl_quantized_modules(self, scheme) -> set[str]: """The diffusers names of the linears the checkpoint quantizes, read from its safetensors header.""" metadata = fetch_checkpoint_metadata(self.bfl_checkpoints[scheme]) if metadata is None: pytest.skip(f"{self.bfl_checkpoints[scheme]} is not a safetensors checkpoint") # one `weight_scale` per quantized linear; the model's single-file mapping renames them to diffusers keys weights = { f"{name[: -len('.weight_scale')]}.weight" for name in metadata.tensors if name.endswith(".weight_scale") } checkpoint = { name: torch.empty(tuple(metadata.tensors[name].shape), device="meta") for name in sorted(weights) } mapping_fn = SINGLE_FILE_LOADABLE_CLASSES[_get_single_file_loadable_mapping_class(self.model_class)][ "checkpoint_mapping_fn" ] converted = mapping_fn(checkpoint=checkpoint, config=None) return {name[: -len(".weight")] for name in converted} def _bfl_quantized_linears(self, model, scheme) -> set[str]: """The checkpoint's quantized linears that the small model has (it has fewer blocks than the checkpoint).""" linears = {name for name, module in model.named_modules() if isinstance(module, torch.nn.Linear)} return self._bfl_quantized_modules(scheme) & linears def _load_unquantized_model(self): torch.manual_seed(0) return self.model_class(**self.get_init_dict()).eval() @torch.no_grad() def _bfl_quantized_state_dict(self, model, scheme): from diffusers.quantizers.bfl.utils import E2M1_MAX, FP8_MAX, quantize_fp8r, quantize_nvfp4 quantized_modules = self._bfl_quantized_linears(model, scheme) assert quantized_modules, "the small model has none of the linears the checkpoint quantizes" # nvfp4 stores a calibrated activation scale per linear: record the input amax on the dummy inputs amax = {} if scheme == "nvfp4": handles = [ model.get_submodule(name).register_forward_pre_hook( lambda module, inputs, name=name: amax.__setitem__(name, inputs[0].abs().amax().item()) ) for name in quantized_modules ] model.to(torch_device)(**self.get_dummy_inputs(), return_dict=False) model.cpu() for handle in handles: handle.remove() state_dict = {} for name, tensor in model.state_dict().items(): module_name = name[: -len(".weight")] if name.endswith(".weight") else None if module_name not in quantized_modules: state_dict[name] = tensor elif scheme == "fp8r": state_dict[name], state_dict[f"{module_name}.weight_scale"] = quantize_fp8r(tensor) else: weight_q, weight_scale, weight_scale_2 = quantize_nvfp4(tensor) state_dict[name] = weight_q state_dict[f"{module_name}.weight_scale"] = weight_scale state_dict[f"{module_name}.weight_scale_2"] = weight_scale_2 # linears the dummy inputs do not reach (e.g. the reference-image blocks) take the largest scale seen input_amax = amax.get(module_name, max(amax.values(), default=1.0)) state_dict[f"{module_name}.input_scale"] = torch.tensor(input_amax / (FP8_MAX * E2M1_MAX)) return state_dict def _create_quantized_model(self, config_kwargs, **extra_kwargs): config_kwargs = dict(config_kwargs) scheme = config_kwargs.pop("scheme") compute_dtype = config_kwargs.get("compute_dtype", torch.bfloat16) model = self._load_unquantized_model().to(compute_dtype) state_dict = self._bfl_quantized_state_dict(model, scheme) kwargs = { "quantization_config": BFLQuantizationConfig(**config_kwargs), "dtype": compute_dtype, "device": str(torch_device), } kwargs.update(extra_kwargs) with tempfile.TemporaryDirectory() as config_dir: model.save_config(config_dir) return self.model_class.from_single_file(state_dict, config=config_dir, **kwargs) def _verify_if_layer_quantized(self, name, module, config_kwargs): cls = self._bfl_linear_cls(config_kwargs["scheme"]) assert isinstance(module, cls), f"Layer {name} is not a {cls.__name__}, got {type(module)}" if config_kwargs["scheme"] == "fp8r": assert module.weight.dtype == torch.float8_e4m3fn, f"{name} weight dtype should be float8_e4m3fn" assert module.weight_scale.shape == (module.out_features,), f"{name} weight_scale is not per row" else: assert module.weight.dtype == torch.uint8, f"{name} weight dtype should be uint8 (packed E2M1)" assert module.weight.shape == (module.out_features, module.in_features // 2), ( f"{name} weight is not packed" ) assert module.weight_scale.dtype == torch.float8_e4m3fn, f"{name} block scales should be float8_e4m3fn" assert module.weight_scale_2.ndim == 0 and module.input_scale.ndim == 0, f"{name} global scales" def _is_module_quantized(self, module, config_kwargs=None): from diffusers.quantizers.bfl.utils import BFLFP8RLinear, BFLNVFP4Linear return isinstance(module, (BFLFP8RLinear, BFLNVFP4Linear)) @is_quantization @is_bfl @require_accelerator @require_accelerate class BFLTesterMixin(BFLConfigMixin, QuantizationTesterMixin): """ Mixin class for testing BFL (`fp8r` / `nvfp4`) quantization on models. Expected class attributes: - model_class: The model class to test Expected methods to be implemented by subclasses: - get_init_dict(): Returns the init kwargs of the small test model - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass Required properties (must be implemented by subclasses): - bfl_checkpoints: Dict mapping a scheme (`"fp8r"`, `"nvfp4"`) to its single-file checkpoint; only the safetensors header is fetched Pytest mark: bfl Use `pytest -m "not bfl"` to skip these tests """ @pytest.mark.parametrize("config_name", list(BFLConfigMixin.BFL_CONFIGS), ids=list(BFLConfigMixin.BFL_CONFIGS)) def test_bfl_quantization_inference(self, config_name): self._test_quantization_inference(BFLConfigMixin.BFL_CONFIGS[config_name]) @pytest.mark.parametrize("config_name", list(BFLConfigMixin.BFL_CONFIGS), ids=list(BFLConfigMixin.BFL_CONFIGS)) def test_bfl_quantization_memory_footprint(self, config_name): self._test_quantization_memory_footprint(BFLConfigMixin.BFL_CONFIGS[config_name]) def _test_quantized_layers(self, config_kwargs): # The exporter quantizes a subset of the linears (embedders, output layers and some blocks stay in bf16), # so the expected set is the one the checkpoint header describes, restricted to the small model's layers. model = self._create_quantized_model(config_kwargs) expected = self._bfl_quantized_linears(model, config_kwargs["scheme"]) quantized = set() for name, module in model.named_modules(): if self._is_module_quantized(module): self._verify_if_layer_quantized(name, module, config_kwargs) quantized.add(name) elif isinstance(module, torch.nn.Linear): assert module.weight.dtype == config_kwargs["compute_dtype"], f"{name} is neither quantized nor bf16" assert quantized, "No quantized layers found in model" assert quantized == expected, f"Quantized layers differ from the checkpoint: {quantized ^ expected}" @pytest.mark.parametrize("config_name", list(BFLConfigMixin.BFL_CONFIGS), ids=list(BFLConfigMixin.BFL_CONFIGS)) def test_bfl_quantized_layers(self, config_name): self._test_quantized_layers(BFLConfigMixin.BFL_CONFIGS[config_name]) @pytest.mark.parametrize("config_name", list(BFLConfigMixin.BFL_CONFIGS), ids=list(BFLConfigMixin.BFL_CONFIGS)) @torch.no_grad() def test_bfl_quantization_matches_unquantized(self, config_name): """The quantized model tracks the model it was quantized from, on the fused kernel or the fallback.""" config_kwargs = BFLConfigMixin.BFL_CONFIGS[config_name] reference = self._load_unquantized_model().to(torch_device, config_kwargs["compute_dtype"]) expected = reference(**self.get_dummy_inputs(), return_dict=False)[0] model = self._create_quantized_model(config_kwargs) model.to(torch_device) output = model(**self.get_dummy_inputs(), return_dict=False)[0] assert output.shape == expected.shape and not torch.isnan(output).any() cosine = torch.nn.functional.cosine_similarity(output.float().flatten(), expected.float().flatten(), dim=0) assert cosine > {"fp8r": 0.99, "nvfp4": 0.9}[config_name], f"cosine similarity {cosine.item():.4f}" @pytest.mark.parametrize("config_name", ["fp8r"], ids=["fp8r"]) def test_bfl_quantization_dtype_assignment(self, config_name): self._test_quantization_dtype_assignment(BFLConfigMixin.BFL_CONFIGS[config_name]) @pytest.mark.parametrize("config_name", ["fp8r"], ids=["fp8r"]) def test_bfl_quantization_lora_inference(self, config_name): self._test_quantization_lora_inference(BFLConfigMixin.BFL_CONFIGS[config_name]) @pytest.mark.parametrize("config_name", list(BFLConfigMixin.BFL_CONFIGS), ids=list(BFLConfigMixin.BFL_CONFIGS)) def test_bfl_keep_modules_in_fp32(self, config_name): # `to_q` is quantized in the checkpoint, so this exercises the quantizer's fp32 exclusion path rather than # a layer the exporter already left in bf16. _keep_in_fp32_modules = self.model_class._keep_in_fp32_modules self.model_class._keep_in_fp32_modules = ["to_q"] try: self._test_keep_modules_in_fp32(BFLConfigMixin.BFL_CONFIGS[config_name]) finally: self.model_class._keep_in_fp32_modules = _keep_in_fp32_modules @pytest.mark.parametrize("config_name", list(BFLConfigMixin.BFL_CONFIGS), ids=list(BFLConfigMixin.BFL_CONFIGS)) def test_bfl_dequantize(self, config_name): self._test_dequantize(BFLConfigMixin.BFL_CONFIGS[config_name]) def test_bfl_rejects_unquantized_checkpoint(self): model = self._load_unquantized_model().to(torch.bfloat16) with tempfile.TemporaryDirectory() as config_dir: model.save_config(config_dir) with pytest.raises(ValueError, match="weight_scale"): self.model_class.from_single_file( model.state_dict(), config=config_dir, quantization_config=BFLQuantizationConfig() ) @is_quantization @is_bfl @is_torch_compile @require_accelerator @require_accelerate class BFLCompileTesterMixin(BFLConfigMixin, QuantizationCompileTesterMixin): """ Mixin class for testing torch.compile with BFL (`fp8r` / `nvfp4`) quantized models. Expected class attributes: - model_class: The model class to test Expected methods to be implemented by subclasses: - get_init_dict(): Returns the init kwargs of the small test model - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass Required properties (must be implemented by subclasses): - bfl_checkpoints: Dict mapping a scheme (`"fp8r"`, `"nvfp4"`) to its single-file checkpoint; only the safetensors header is fetched Pytest mark: bfl Use `pytest -m "not bfl"` to skip these tests """ @pytest.mark.parametrize("config_name", list(BFLConfigMixin.BFL_CONFIGS), ids=list(BFLConfigMixin.BFL_CONFIGS)) def test_bfl_torch_compile(self, config_name): if ( config_name == "nvfp4" and torch_device == "cuda" and torch.cuda.get_device_capability()[0] >= 10 and is_flashinfer_available() ): # Dynamo cannot trace FlashInfer's Python wrappers (they take a thread lock), and with graph breaks around # them the remaining fragments recompile between calls. Expected to pass once the fused nvfp4 path wraps # the FlashInfer calls as custom ops. pytest.xfail("the fused nvfp4 path calls FlashInfer, which torch.compile cannot trace") self._test_torch_compile(BFLConfigMixin.BFL_CONFIGS[config_name]) @pytest.mark.parametrize("config_name", ["fp8r"], ids=["fp8r"]) def test_bfl_torch_compile_with_group_offload(self, config_name): self._test_torch_compile_with_group_offload(BFLConfigMixin.BFL_CONFIGS[config_name])