From b695febc16287b2b8cff8f89a0120ea71bbdd29b Mon Sep 17 00:00:00 2001 From: "sayak@huggingface.co" Date: Mon, 10 Aug 2026 11:22:10 +0000 Subject: [PATCH 1/6] fix(torchao): route root-level tensors around safetensors reconstruction Models with parameters at the root of the module tree (e.g. Wan's `scale_shift_table`) crashed torchao's `unflatten_tensor_state_dict` when loading serialized checkpoints, since flattened tensor names are assumed to carry a module prefix. Filter such tensors and their metadata entries out of the reconstruction and merge them back unchanged. Co-Authored-By: Claude Fable 5 --- .../quantizers/torchao/torchao_quantizer.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/src/diffusers/quantizers/torchao/torchao_quantizer.py b/src/diffusers/quantizers/torchao/torchao_quantizer.py index 503e452c20b1..91aab04975aa 100644 --- a/src/diffusers/quantizers/torchao/torchao_quantizer.py +++ b/src/diffusers/quantizers/torchao/torchao_quantizer.py @@ -302,9 +302,19 @@ def maybe_update_state_dict(self, state_dict: dict[str, Any]) -> dict[str, Any]: return state_dict merged_state_dict = {**self._pending_flattened_state_dict, **state_dict} + # Tensors at the model root (e.g. Wan's `scale_shift_table`) have no module prefix and are never + # flattened tensor-subclass parts; torchao's unflatten helper cannot parse their names, so route + # them (and their metadata entries) around the reconstruction. + root_tensors = {k: v for k, v in merged_state_dict.items() if "." not in k} + merged_state_dict = {k: v for k, v in merged_state_dict.items() if "." in k} + metadata = self._metadata + tensor_names = json.loads(metadata["tensor_names"]) + if any("." not in name for name in tensor_names): + metadata = {**metadata, "tensor_names": json.dumps([name for name in tensor_names if "." in name])} reconstructed_state_dict, self._pending_flattened_state_dict = unflatten_tensor_state_dict( - merged_state_dict, self._metadata + merged_state_dict, metadata ) + reconstructed_state_dict.update(root_tensors) return reconstructed_state_dict From bb35ec9f4fd51afae81ee2f392ef5dc614d89ae4 Mon Sep 17 00:00:00 2001 From: "sayak@huggingface.co" Date: Mon, 10 Aug 2026 11:22:10 +0000 Subject: [PATCH 2/6] fix(bnb): defer 8-bit weights split from their SCB stats across shards Sharded serialization can place an 8-bit weight and its `SCB` statistics in different shard files, in which case the shard-by-shard loader failed with "Missing quantization component `SCB`". Hold the incomplete half of the pair back until its counterpart arrives with a later shard, mirroring the torchao pending mechanism, and disable parallel shard loading for prequantized 8-bit checkpoints. Co-Authored-By: Claude Fable 5 --- .../quantizers/bitsandbytes/bnb_quantizer.py | 33 +++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/src/diffusers/quantizers/bitsandbytes/bnb_quantizer.py b/src/diffusers/quantizers/bitsandbytes/bnb_quantizer.py index 7c5dcbc73d5f..63d2680b3d33 100644 --- a/src/diffusers/quantizers/bitsandbytes/bnb_quantizer.py +++ b/src/diffusers/quantizers/bitsandbytes/bnb_quantizer.py @@ -345,6 +345,39 @@ def __init__(self, quantization_config, **kwargs): if self.quantization_config.llm_int8_skip_modules is not None: self.modules_to_not_convert = self.quantization_config.llm_int8_skip_modules + self._checkpoint_keys = set() + self._pending_quantized_state = {} + + def maybe_update_loaded_keys(self, loaded_keys: list[str], checkpoint_files: list[str]) -> list[str]: + self._checkpoint_keys = set(loaded_keys) + return loaded_keys + + def maybe_update_state_dict(self, state_dict: dict[str, Any]) -> dict[str, Any]: + if not self.pre_quantized: + return state_dict + + # A sharded checkpoint can split an 8-bit weight from its `SCB` statistics, which must be + # materialized together. Hold the incomplete half back until its counterpart arrives with a + # later shard. + merged = {**self._pending_quantized_state, **state_dict} + pending = {} + for name in list(merged.keys()): + if name.endswith(".weight"): + partner = name[: -len("weight")] + "SCB" + elif name.endswith(".SCB"): + partner = name[: -len("SCB")] + "weight" + else: + continue + if partner in self._checkpoint_keys and partner not in merged: + pending[name] = merged.pop(name) + self._pending_quantized_state = pending + return merged + + @property + def supports_parallel_loading(self) -> bool: + # Deferred SCB reconstruction carries incomplete weight/SCB pairs from one shard to the next. + return not self.pre_quantized + def validate_environment(self, *args, **kwargs): if not (torch.cuda.is_available() or torch.xpu.is_available()): raise RuntimeError("No GPU found. A GPU is needed for quantization.") From d4a1ad65aa4d9e614c6de592f25d1f00af041fde Mon Sep 17 00:00:00 2001 From: "sayak@huggingface.co" Date: Mon, 10 Aug 2026 09:32:15 +0000 Subject: [PATCH 3/6] [tests] extend model-level quantization tester mixins Migrate remaining model-level coverage from tests/quantization into the tester mixins so it runs for every wired model: - base: buffer-placement assertions in the device-map test, and an opt-in sharded-serialization test enabled by setting `sharded_serialization_config` - bnb: serialization across all configs (sharded included), dtype assignment and adapter training for 8-bit, device moves preserving the memory footprint, corrupted-state-dict loading error, and a fixed modules-to-not-convert test (BitsAndBytesConfig only exposes llm_int8_skip_modules; the old test passed an unsupported kwarg and only survived by being skipped) - torchao: custom device maps with cpu/disk offload, generalized from the Flux-specific test - gguf: the diffusers-format single-file loading path, wired into the Flux model tests Co-Authored-By: Claude Fable 5 --- tests/models/testing_utils/quantization.py | 180 ++++++++++++++++-- .../test_models_transformer_flux.py | 15 +- 2 files changed, 173 insertions(+), 22 deletions(-) diff --git a/tests/models/testing_utils/quantization.py b/tests/models/testing_utils/quantization.py index 38dc4ded4b68..d7c117c13878 100644 --- a/tests/models/testing_utils/quantization.py +++ b/tests/models/testing_utils/quantization.py @@ -16,6 +16,7 @@ import gc import pytest +import safetensors.torch import torch from diffusers import ( @@ -145,14 +146,15 @@ def _create_quantized_model(self, config_kwargs, **extra_kwargs): 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): + 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. + 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, {}) + self._verify_if_layer_quantized("", module, config_kwargs or {}) return True except (AssertionError, AttributeError): return False @@ -243,6 +245,14 @@ def _test_quantization_lora_inference(self, config_kwargs): 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): """ @@ -251,7 +261,10 @@ def _test_quantization_serialization(self, config_kwargs, tmp_path, max_shard_si 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 + 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) @@ -259,6 +272,9 @@ def _test_quantization_serialization(self, config_kwargs, tmp_path, max_shard_si 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 @@ -360,17 +376,22 @@ def _test_keep_modules_in_fp32(self, config_kwargs): f"Module {name} should be FP32 but is {module.weight.dtype}" ) - def _test_quantization_modules_to_not_convert(self, config_kwargs, modules_to_not_convert): + @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 modules_to_not_convert + # Create config with the exclusion list config_kwargs_with_exclusion = config_kwargs.copy() - config_kwargs_with_exclusion["modules_to_not_convert"] = modules_to_not_convert + config_kwargs_with_exclusion[exclusion_key] = modules_to_not_convert model_with_exclusion = self._create_quantized_model(config_kwargs_with_exclusion) @@ -382,7 +403,7 @@ def _test_quantization_modules_to_not_convert(self, config_kwargs, modules_to_no 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), ( + assert not self._is_module_quantized(module, config_kwargs), ( f"Module {name} should not be quantized but was found to be quantized" ) @@ -394,12 +415,19 @@ def _test_quantization_modules_to_not_convert(self, config_kwargs, modules_to_no 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): + 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) @@ -423,6 +451,15 @@ 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 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" @@ -461,7 +498,9 @@ def _test_dequantize(self, config_kwargs): for name, module in model.named_modules(): if isinstance(module, torch.nn.Linear): - assert not self._is_module_quantized(module), f"Module {name} is still quantized after dequantize()" + 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] @@ -511,8 +550,9 @@ def _test_quantization_training(self, config_kwargs): # Step 3: run forward and backward pass inputs = self.get_dummy_inputs() - # Use bfloat16 on XPU to avoid gradient underflow with quantized layers - autocast_dtype = torch.bfloat16 if torch_device == "xpu" else torch.float16 + # 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() @@ -561,6 +601,8 @@ class BitsAndBytesConfigMixin: "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() @@ -612,7 +654,7 @@ def test_bnb_quantization_num_parameters(self, config_name): ids=list(BitsAndBytesConfigMixin.BNB_CONFIGS.keys()), ) def test_bnb_quantization_memory_footprint(self, config_name): - expected = BitsAndBytesConfigMixin.BNB_EXPECTED_MEMORY_REDUCTIONS.get(config_name, 1.2) + 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 ) @@ -625,15 +667,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) @@ -659,16 +718,39 @@ def test_bnb_original_dtype(self): 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.""" + @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["4bit_nf4"], modules_to_exclude + 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.""" @@ -678,9 +760,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", @@ -1012,10 +1095,65 @@ def test_torchao_modules_to_not_convert(self): 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"]) @@ -1593,7 +1731,7 @@ class NunchakuLiteTesterMixin(NunchakuLiteConfigMixin, QuantizationTesterMixin): def test_nunchaku_lite_quantization_inference(self): self._test_quantization_inference(self.config_dict) - def _is_module_quantized(self, module): + def _is_module_quantized(self, module, config_kwargs=None): from diffusers.quantizers.nunchaku.utils import AWQW4A16Linear, SVDQW4A4Linear return isinstance(module, (SVDQW4A4Linear, AWQW4A16Linear)) diff --git a/tests/models/transformers/test_models_transformer_flux.py b/tests/models/transformers/test_models_transformer_flux.py index 719429526945..b233bbe9fd1d 100644 --- a/tests/models/transformers/test_models_transformer_flux.py +++ b/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 @@ -410,6 +412,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.""" From 5e149c46e869073a33bffdc4c326833c7e36622d Mon Sep 17 00:00:00 2001 From: "sayak@huggingface.co" Date: Mon, 10 Aug 2026 11:22:11 +0000 Subject: [PATCH 4/6] [tests] fix model-level quantization test wiring - SD3.5: the quantized testers reused the random-init dummy inputs (4 latent channels, fp32) while the tiny Hub checkpoint has in_channels=8 and the quantizers load the model in half precision; give them matching inputs and relax the 4-bit memory expectation for the tiny checkpoint. - QwenImage / Flux2: the quantized testers had no Hub checkpoint wired at all, so every test errored; point them at hf-internal-testing/tiny-qwenimage-pipe and tiny-flux2 with matching inputs. - NucleusMoE: no tiny checkpoint exists on the Hub yet; comment the testers out like the LTX ones. Co-Authored-By: Claude Fable 5 --- .../test_models_transformer_flux.py | 12 +++++ .../test_models_transformer_flux2.py | 38 +++++++++++++-- ...est_models_transformer_nucleusmoe_image.py | 12 ++--- .../test_models_transformer_qwenimage.py | 32 ++++++++++++- .../test_models_transformer_sd3.py | 48 ++++++++++++++++++- 5 files changed, 129 insertions(+), 13 deletions(-) diff --git a/tests/models/transformers/test_models_transformer_flux.py b/tests/models/transformers/test_models_transformer_flux.py index b233bbe9fd1d..1add8f81f6b1 100644 --- a/tests/models/transformers/test_models_transformer_flux.py +++ b/tests/models/transformers/test_models_transformer_flux.py @@ -362,6 +362,12 @@ class TestFluxTransformerBitsAndBytes(FluxTransformerTesterConfig, BitsAndBytesT def torch_dtype(self): return torch.float16 + def get_dummy_inputs(self): + """Override to build inputs in the quantizer compute dtype (excluded/unquantized linears + do strict-dtype matmuls).""" + inputs = super().get_dummy_inputs() + return {k: v.to(self.torch_dtype) if torch.is_floating_point(v) else v for k, v in inputs.items()} + class TestFluxTransformerQuanto(FluxTransformerTesterConfig, QuantoTesterMixin): """Quanto quantization tests for Flux Transformer.""" @@ -384,6 +390,12 @@ class TestFluxTransformerTorchAo(FluxTransformerTesterConfig, TorchAoTesterMixin def torch_dtype(self): return torch.bfloat16 + def get_dummy_inputs(self): + """Override to build inputs in the quantizer compute dtype (excluded/unquantized linears + do strict-dtype matmuls).""" + inputs = super().get_dummy_inputs() + return {k: v.to(self.torch_dtype) if torch.is_floating_point(v) else v for k, v in inputs.items()} + class TestFluxTransformerGGUF(FluxTransformerTesterConfig, GGUFTesterMixin): @property diff --git a/tests/models/transformers/test_models_transformer_flux2.py b/tests/models/transformers/test_models_transformer_flux2.py index 9546fdb5d969..b279b57dcc6e 100644 --- a/tests/models/transformers/test_models_transformer_flux2.py +++ b/tests/models/transformers/test_models_transformer_flux2.py @@ -251,13 +251,41 @@ def get_dummy_inputs(self, height: int = 4, width: int = 4) -> dict[str, torch.T } -class TestFlux2TransformerBitsAndBytes(Flux2TransformerTesterConfig, BitsAndBytesTesterMixin): +class Flux2TransformerQuantTesterConfig(Flux2TransformerTesterConfig): + """Shared config for quantized Flux2 Transformer tests (loads the tiny Hub checkpoint).""" + + @property + def pretrained_model_name_or_path(self): + return "hf-internal-testing/tiny-flux2" + + @property + def pretrained_model_kwargs(self): + return {"subfolder": "transformer"} + + def get_dummy_inputs(self, height: int = 4, width: int = 4, batch_size: int = 1) -> dict[str, torch.Tensor]: + """Override to match the tiny Hub checkpoint (joint_attention_dim=16) and the quantizer compute dtype.""" + inputs = super().get_dummy_inputs(height=height, width=width, batch_size=batch_size) + inputs["encoder_hidden_states"] = randn_tensor( + (batch_size, inputs["encoder_hidden_states"].shape[1], 16), generator=self.generator, device=torch_device + ) + return {k: v.to(self.torch_dtype) if torch.is_floating_point(v) else v for k, v in inputs.items()} + + +class TestFlux2TransformerBitsAndBytes(Flux2TransformerQuantTesterConfig, BitsAndBytesTesterMixin): """BitsAndBytes quantization tests for Flux2 Transformer.""" + @property + def torch_dtype(self): + return torch.float16 + -class TestFlux2TransformerTorchAo(Flux2TransformerTesterConfig, TorchAoTesterMixin): +class TestFlux2TransformerTorchAo(Flux2TransformerQuantTesterConfig, TorchAoTesterMixin): """TorchAO quantization tests for Flux2 Transformer.""" + @property + def torch_dtype(self): + return torch.bfloat16 + class TestFlux2TransformerGGUF(Flux2TransformerTesterConfig, GGUFTesterMixin): """GGUF quantization tests for Flux2 Transformer.""" @@ -315,9 +343,13 @@ def get_dummy_inputs(self): } -class TestFlux2TransformerTorchAoCompile(Flux2TransformerTesterConfig, TorchAoCompileTesterMixin): +class TestFlux2TransformerTorchAoCompile(Flux2TransformerQuantTesterConfig, TorchAoCompileTesterMixin): """TorchAO + compile tests for Flux2 Transformer.""" + @property + def torch_dtype(self): + return torch.bfloat16 + class TestFlux2TransformerGGUFCompile(Flux2TransformerTesterConfig, GGUFCompileTesterMixin): """GGUF + compile tests for Flux2 Transformer.""" diff --git a/tests/models/transformers/test_models_transformer_nucleusmoe_image.py b/tests/models/transformers/test_models_transformer_nucleusmoe_image.py index f4ecb9a2c62a..9b7e70f345cf 100644 --- a/tests/models/transformers/test_models_transformer_nucleusmoe_image.py +++ b/tests/models/transformers/test_models_transformer_nucleusmoe_image.py @@ -22,12 +22,10 @@ from ..testing_utils import ( AttentionTesterMixin, BaseModelTesterConfig, - BitsAndBytesTesterMixin, LoraHotSwappingForModelTesterMixin, LoraTesterMixin, MemoryTesterMixin, ModelTesterMixin, - TorchAoTesterMixin, TorchCompileTesterMixin, TrainingTesterMixin, ) @@ -212,9 +210,11 @@ def get_dummy_inputs(self, height: int = 4, width: int = 4) -> dict: } -class TestNucleusMoEImageTransformerBitsAndBytes(NucleusMoEImageTransformerTesterConfig, BitsAndBytesTesterMixin): - """BitsAndBytes quantization tests for NucleusMoE Image Transformer.""" +# TODO: Add pretrained_model_name_or_path once a tiny NucleusMoE model is available on the Hub +# class TestNucleusMoEImageTransformerBitsAndBytes(NucleusMoEImageTransformerTesterConfig, BitsAndBytesTesterMixin): +# """BitsAndBytes quantization tests for NucleusMoE Image Transformer.""" -class TestNucleusMoEImageTransformerTorchAo(NucleusMoEImageTransformerTesterConfig, TorchAoTesterMixin): - """TorchAO quantization tests for NucleusMoE Image Transformer.""" +# TODO: Add pretrained_model_name_or_path once a tiny NucleusMoE model is available on the Hub +# class TestNucleusMoEImageTransformerTorchAo(NucleusMoEImageTransformerTesterConfig, TorchAoTesterMixin): +# """TorchAO quantization tests for NucleusMoE Image Transformer.""" diff --git a/tests/models/transformers/test_models_transformer_qwenimage.py b/tests/models/transformers/test_models_transformer_qwenimage.py index bed2bda4064f..1a207d3f4068 100644 --- a/tests/models/transformers/test_models_transformer_qwenimage.py +++ b/tests/models/transformers/test_models_transformer_qwenimage.py @@ -430,9 +430,37 @@ def test_torch_compile_with_and_without_mask(self): assert not torch.allclose(output_no_mask.sample, output_with_padding.sample, atol=1e-3) -class TestQwenImageTransformerBitsAndBytes(QwenImageTransformerTesterConfig, BitsAndBytesTesterMixin): +class QwenImageTransformerQuantTesterConfig(QwenImageTransformerTesterConfig): + """Shared config for quantized QwenImage Transformer tests (loads the tiny Hub checkpoint).""" + + @property + def pretrained_model_name_or_path(self): + return "hf-internal-testing/tiny-qwenimage-pipe" + + @property + def pretrained_model_kwargs(self): + return {"subfolder": "transformer"} + + def get_dummy_inputs(self, batch_size: int = 1) -> dict[str, torch.Tensor]: + """Override to match the compute dtype the quantizer loads the model in.""" + inputs = super().get_dummy_inputs(batch_size=batch_size) + return { + k: v.to(self.torch_dtype) if torch.is_tensor(v) and torch.is_floating_point(v) else v + for k, v in inputs.items() + } + + +class TestQwenImageTransformerBitsAndBytes(QwenImageTransformerQuantTesterConfig, BitsAndBytesTesterMixin): """BitsAndBytes quantization tests for QwenImage Transformer.""" + @property + def torch_dtype(self): + return torch.float16 -class TestQwenImageTransformerTorchAo(QwenImageTransformerTesterConfig, TorchAoTesterMixin): + +class TestQwenImageTransformerTorchAo(QwenImageTransformerQuantTesterConfig, TorchAoTesterMixin): """TorchAO quantization tests for QwenImage Transformer.""" + + @property + def torch_dtype(self): + return torch.bfloat16 diff --git a/tests/models/transformers/test_models_transformer_sd3.py b/tests/models/transformers/test_models_transformer_sd3.py index e38c7853a613..2f1d5fff6e31 100644 --- a/tests/models/transformers/test_models_transformer_sd3.py +++ b/tests/models/transformers/test_models_transformer_sd3.py @@ -218,9 +218,53 @@ class TestSD35TransformerCompile(SD35TransformerTesterConfig, TorchCompileTester pass -class TestSD35TransformerBitsAndBytes(SD35TransformerTesterConfig, BitsAndBytesTesterMixin): +class SD35TransformerQuantTesterConfig(SD35TransformerTesterConfig): + """Shared config for quantized SD3.5 Transformer tests (matches the tiny Hub checkpoint).""" + + def get_dummy_inputs(self, batch_size: int = 2) -> dict[str, torch.Tensor]: + """Override to match the tiny Hub checkpoint (in_channels=8) and the quantizer compute dtype.""" + num_channels = 8 + height = width = embedding_dim = 32 + pooled_embedding_dim = embedding_dim * 2 + sequence_length = 154 + + return { + "hidden_states": randn_tensor( + (batch_size, num_channels, height, width), + generator=self.generator, + device=torch_device, + dtype=self.torch_dtype, + ), + "encoder_hidden_states": randn_tensor( + (batch_size, sequence_length, embedding_dim), + generator=self.generator, + device=torch_device, + dtype=self.torch_dtype, + ), + "pooled_projections": randn_tensor( + (batch_size, pooled_embedding_dim), + generator=self.generator, + device=torch_device, + dtype=self.torch_dtype, + ), + "timestep": torch.randint(0, 1000, size=(batch_size,), generator=self.generator).to(torch_device), + } + + +class TestSD35TransformerBitsAndBytes(SD35TransformerQuantTesterConfig, BitsAndBytesTesterMixin): """BitsAndBytes quantization tests for SD3.5 Transformer.""" + # The tiny SD3.5 transformer has a smaller linear fraction than the 4-bit default expectation of 3.0x. + BNB_EXPECTED_MEMORY_REDUCTIONS = {"4bit_nf4": 2.5, "4bit_fp4": 2.5, "8bit": 1.5} -class TestSD35TransformerTorchAo(SD35TransformerTesterConfig, TorchAoTesterMixin): + @property + def torch_dtype(self): + return torch.float16 + + +class TestSD35TransformerTorchAo(SD35TransformerQuantTesterConfig, TorchAoTesterMixin): """TorchAO quantization tests for SD3.5 Transformer.""" + + @property + def torch_dtype(self): + return torch.bfloat16 From 24183aba41585c61c2c4aa984e90579407377d36 Mon Sep 17 00:00:00 2001 From: "sayak@huggingface.co" Date: Mon, 10 Aug 2026 09:32:15 +0000 Subject: [PATCH 5/6] [tests] split tests/quantization into pipeline-level and backend-level tiers Pipeline-level quantization tests (pipeline quality slices, cpu offload, LoRA loading, compile, PipelineQuantizationConfig) move to tests/pipelines/testing_utils/quantization.py, marked per backend so the nightly CI can select them with `pytest -m`. tests/quantization keeps only backend-level tests that fit neither tier: config validation, utility warnings, and GGUF CUDA kernel correctness. Tests already covered by the model-level mixins are dropped. Co-Authored-By: Claude Fable 5 --- tests/pipelines/testing_utils/quantization.py | 1695 +++++++++++++++++ tests/quantization/bnb/README.md | 41 - tests/quantization/bnb/test_4bit.py | 633 +----- tests/quantization/bnb/test_mixed_int8.py | 759 +------- tests/quantization/gguf/test_gguf.py | 699 +------ tests/quantization/modelopt/__init__.py | 0 tests/quantization/modelopt/test_modelopt.py | 343 ---- tests/quantization/quanto/__init__.py | 0 tests/quantization/quanto/test_quanto.py | 176 -- .../test_pipeline_level_quantization.py | 312 --- .../quantization/test_torch_compile_utils.py | 66 - tests/quantization/torchao/README.md | 50 - tests/quantization/torchao/test_torchao.py | 689 +------ tests/quantization/utils.py | 45 - 14 files changed, 1741 insertions(+), 3767 deletions(-) create mode 100644 tests/pipelines/testing_utils/quantization.py delete mode 100644 tests/quantization/bnb/README.md delete mode 100644 tests/quantization/modelopt/__init__.py delete mode 100644 tests/quantization/modelopt/test_modelopt.py delete mode 100644 tests/quantization/quanto/__init__.py delete mode 100644 tests/quantization/quanto/test_quanto.py delete mode 100644 tests/quantization/test_pipeline_level_quantization.py delete mode 100644 tests/quantization/test_torch_compile_utils.py delete mode 100644 tests/quantization/torchao/README.md delete mode 100644 tests/quantization/utils.py diff --git a/tests/pipelines/testing_utils/quantization.py b/tests/pipelines/testing_utils/quantization.py new file mode 100644 index 000000000000..547751b9972f --- /dev/null +++ b/tests/pipelines/testing_utils/quantization.py @@ -0,0 +1,1695 @@ +# coding=utf-8 +# Copyright 2026 The HuggingFace Team 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. +"""Pipeline-level quantization tests. + +Model-level quantization tests live in `tests/models/testing_utils/quantization.py` and are wired +into the individual model test files. Backend-level tests (config validation, loading error paths, +kernels) live in `tests/quantization/`. This module only covers behavior that needs a pipeline. + +The module name intentionally does not match pytest's `test_*.py` discovery pattern: these tests +only run when this file is passed to pytest explicitly, as the nightly quantization CI jobs do. +Every class is marked with the `quantization` marker plus its backend marker (`bitsandbytes`, +`torchao`, `gguf`, `modelopt`) so CI can select per-backend subsets with `pytest -m`. +""" + +import gc +import json +import tempfile + +import numpy as np +import pytest +import torch +from huggingface_hub import hf_hub_download +from parameterized import parameterized +from PIL import Image + +from diffusers import ( + AuraFlowPipeline, + AuraFlowTransformer2DModel, + AutoencoderKL, + BitsAndBytesConfig, + DiffusionPipeline, + FlowMatchEulerDiscreteScheduler, + FluxControlPipeline, + FluxPipeline, + FluxTransformer2DModel, + GGUFQuantizationConfig, + NVIDIAModelOptConfig, + QuantoConfig, + SD3Transformer2DModel, + StableDiffusion3Pipeline, + TorchAoConfig, +) +from diffusers.quantizers import PipelineQuantizationConfig +from diffusers.utils import is_accelerate_version, load_image, logging + +from ...testing_utils import ( + CaptureLogger, + Expectations, + backend_empty_cache, + backend_reset_peak_memory_stats, + backend_synchronize, + enable_full_determinism, + is_bitsandbytes, + is_bitsandbytes_available, + is_gguf, + is_gguf_available, + is_modelopt, + is_quantization, + is_torchao, + is_torchao_available, + is_transformers_available, + nightly, + numpy_cosine_similarity_distance, + require_accelerate, + require_big_accelerator, + require_bitsandbytes_version_greater, + require_gguf_version_greater_or_equal, + require_modelopt_version_greater_or_equal, + require_peft_backend, + require_peft_version_greater, + require_quanto, + require_torch, + require_torch_accelerator, + require_torch_version_greater, + require_torch_version_greater_equal, + require_torchao_version_greater_or_equal, + require_transformers_version_greater, + slow, + torch_device, +) + + +if is_transformers_available(): + from transformers import AutoTokenizer, CLIPTextModel, CLIPTokenizer, T5EncoderModel + from transformers import BitsAndBytesConfig as TranBitsAndBytesConfig +else: + TranBitsAndBytesConfig = None + +if is_bitsandbytes_available(): + pass + +if is_torchao_available(): + from torchao.quantization import ( + Float8WeightOnlyConfig, + Int4WeightOnlyConfig, + Int8DynamicActivationInt8WeightConfig, + Int8DynamicActivationIntxWeightConfig, + Int8Tensor, + Int8WeightOnlyConfig, + IntxWeightOnlyConfig, + ) + from torchao.utils import TorchAOBaseTensor + +if is_gguf_available(): + pass + + +enable_full_determinism() + + +# ======================== Shared compile base ======================== + + +@is_quantization +@require_torch_accelerator +@slow +class QuantCompileTests: + @property + def quantization_config(self): + raise NotImplementedError( + "This property should be implemented in the subclass to return the appropriate quantization config." + ) + + @pytest.fixture(autouse=True) + def _cleanup(self): + gc.collect() + backend_empty_cache(torch_device) + torch.compiler.reset() + yield + gc.collect() + backend_empty_cache(torch_device) + torch.compiler.reset() + + def _init_pipeline(self, quantization_config, torch_dtype): + pipe = DiffusionPipeline.from_pretrained( + "stabilityai/stable-diffusion-3-medium-diffusers", + quantization_config=quantization_config, + torch_dtype=torch_dtype, + ) + return pipe + + def _test_torch_compile_with_cpu_offload(self, torch_dtype=torch.bfloat16): + pipe = self._init_pipeline(self.quantization_config, torch_dtype) + pipe.enable_model_cpu_offload() + # regional compilation is better for offloading. + # see: https://pytorch.org/blog/torch-compile-and-diffusers-a-hands-on-guide-to-peak-performance/ + if getattr(pipe.transformer, "_repeated_blocks"): + pipe.transformer.compile_repeated_blocks(fullgraph=True) + else: + pipe.transformer.compile() + + # small resolutions to ensure speedy execution. + pipe("a dog", num_inference_steps=2, max_sequence_length=16, height=256, width=256) + + def test_torch_compile_with_cpu_offload(self): + self._test_torch_compile_with_cpu_offload() + + +# ======================== PipelineQuantizationConfig ======================== + + +@is_quantization +@require_bitsandbytes_version_greater("0.43.2") +@require_quanto +@require_accelerate +@require_torch +@require_torch_accelerator +@slow +class TestPipelineQuantization: + model_name = "hf-internal-testing/tiny-flux-pipe" + prompt = "a beautiful sunset amidst the mountains." + num_inference_steps = 10 + seed = 0 + + def test_quant_config_set_correctly_through_kwargs(self): + components_to_quantize = ["transformer", "text_encoder_2"] + quant_config = PipelineQuantizationConfig( + quant_backend="bitsandbytes_4bit", + quant_kwargs={ + "load_in_4bit": True, + "bnb_4bit_quant_type": "nf4", + "bnb_4bit_compute_dtype": torch.bfloat16, + }, + components_to_quantize=components_to_quantize, + ) + pipe = DiffusionPipeline.from_pretrained( + self.model_name, + quantization_config=quant_config, + torch_dtype=torch.bfloat16, + ).to(torch_device) + for name, component in pipe.components.items(): + if name in components_to_quantize: + assert getattr(component.config, "quantization_config", None) is not None + quantization_config = component.config.quantization_config + assert quantization_config.load_in_4bit + assert quantization_config.quant_method == "bitsandbytes" + + _ = pipe(self.prompt, num_inference_steps=self.num_inference_steps) + + def test_quant_config_set_correctly_through_granular(self): + quant_config = PipelineQuantizationConfig( + quant_mapping={ + "transformer": QuantoConfig(weights_dtype="int8"), + "text_encoder_2": TranBitsAndBytesConfig(load_in_4bit=True, compute_dtype=torch.bfloat16), + } + ) + components_to_quantize = list(quant_config.quant_mapping.keys()) + pipe = DiffusionPipeline.from_pretrained( + self.model_name, + quantization_config=quant_config, + torch_dtype=torch.bfloat16, + ).to(torch_device) + for name, component in pipe.components.items(): + if name in components_to_quantize: + assert getattr(component.config, "quantization_config", None) is not None + quantization_config = component.config.quantization_config + + if name == "text_encoder_2": + assert quantization_config.load_in_4bit + assert quantization_config.quant_method == "bitsandbytes" + else: + assert quantization_config.quant_method == "quanto" + + _ = pipe(self.prompt, num_inference_steps=self.num_inference_steps) + + def test_raises_error_for_invalid_config(self): + with pytest.raises(ValueError) as err_context: + _ = PipelineQuantizationConfig( + quant_mapping={ + "transformer": QuantoConfig(weights_dtype="int8"), + "text_encoder_2": TranBitsAndBytesConfig(load_in_4bit=True, compute_dtype=torch.bfloat16), + }, + quant_backend="bitsandbytes_4bit", + ) + + assert ( + str(err_context.value) == "Both `quant_backend` and `quant_mapping` cannot be specified at the same time." + ) + + def test_validation_for_kwargs(self): + components_to_quantize = ["transformer", "text_encoder_2"] + with pytest.raises(ValueError) as err_context: + _ = PipelineQuantizationConfig( + quant_backend="quanto", + quant_kwargs={"weights_dtype": "int8"}, + components_to_quantize=components_to_quantize, + ) + + assert "The signatures of the __init__ methods of the quantization config classes" in str(err_context.value) + + def test_raises_error_for_wrong_config_class(self): + quant_config = { + "transformer": QuantoConfig(weights_dtype="int8"), + "text_encoder_2": TranBitsAndBytesConfig(load_in_4bit=True, compute_dtype=torch.bfloat16), + } + with pytest.raises(ValueError) as err_context: + _ = DiffusionPipeline.from_pretrained( + self.model_name, + quantization_config=quant_config, + torch_dtype=torch.bfloat16, + ) + assert str(err_context.value) == "`quantization_config` must be an instance of `PipelineQuantizationConfig`." + + def test_validation_for_mapping(self): + with pytest.raises(ValueError) as err_context: + _ = PipelineQuantizationConfig( + quant_mapping={ + "transformer": DiffusionPipeline(), + "text_encoder_2": TranBitsAndBytesConfig(load_in_4bit=True, compute_dtype=torch.bfloat16), + } + ) + + assert "Provided config for module_name=transformer could not be found" in str(err_context.value) + + def test_saving_loading(self): + quant_config = PipelineQuantizationConfig( + quant_mapping={ + "transformer": QuantoConfig(weights_dtype="int8"), + "text_encoder_2": TranBitsAndBytesConfig(load_in_4bit=True, compute_dtype=torch.bfloat16), + } + ) + components_to_quantize = list(quant_config.quant_mapping.keys()) + pipe = DiffusionPipeline.from_pretrained( + self.model_name, + quantization_config=quant_config, + torch_dtype=torch.bfloat16, + ).to(torch_device) + + pipe_inputs = {"prompt": self.prompt, "num_inference_steps": self.num_inference_steps, "output_type": "latent"} + output_1 = pipe(**pipe_inputs, generator=torch.manual_seed(self.seed)).images + + with tempfile.TemporaryDirectory() as tmpdir: + pipe.save_pretrained(tmpdir) + loaded_pipe = DiffusionPipeline.from_pretrained(tmpdir, torch_dtype=torch.bfloat16).to(torch_device) + for name, component in loaded_pipe.components.items(): + if name in components_to_quantize: + assert getattr(component.config, "quantization_config", None) is not None + quantization_config = component.config.quantization_config + + if name == "text_encoder_2": + assert quantization_config.load_in_4bit + assert quantization_config.quant_method == "bitsandbytes" + else: + assert quantization_config.quant_method == "quanto" + + output_2 = loaded_pipe(**pipe_inputs, generator=torch.manual_seed(self.seed)).images + + assert torch.allclose(output_1, output_2) + + @parameterized.expand(["quant_kwargs", "quant_mapping"]) + def test_warn_invalid_component(self, method): + invalid_component = "foo" + if method == "quant_kwargs": + components_to_quantize = ["transformer", invalid_component] + quant_config = PipelineQuantizationConfig( + quant_backend="bitsandbytes_8bit", + quant_kwargs={"load_in_8bit": True}, + components_to_quantize=components_to_quantize, + ) + else: + quant_config = PipelineQuantizationConfig( + quant_mapping={ + "transformer": QuantoConfig("int8"), + invalid_component: TranBitsAndBytesConfig(load_in_8bit=True), + } + ) + + logger = logging.get_logger("diffusers.pipelines.pipeline_loading_utils") + logger.setLevel(logging.WARNING) + with CaptureLogger(logger) as cap_logger: + _ = DiffusionPipeline.from_pretrained( + self.model_name, + quantization_config=quant_config, + torch_dtype=torch.bfloat16, + ) + assert invalid_component in cap_logger.out + + @parameterized.expand(["quant_kwargs", "quant_mapping"]) + def test_no_quantization_for_all_invalid_components(self, method): + invalid_component = "foo" + if method == "quant_kwargs": + components_to_quantize = [invalid_component] + quant_config = PipelineQuantizationConfig( + quant_backend="bitsandbytes_8bit", + quant_kwargs={"load_in_8bit": True}, + components_to_quantize=components_to_quantize, + ) + else: + quant_config = PipelineQuantizationConfig( + quant_mapping={invalid_component: TranBitsAndBytesConfig(load_in_8bit=True)} + ) + + pipe = DiffusionPipeline.from_pretrained( + self.model_name, + quantization_config=quant_config, + torch_dtype=torch.bfloat16, + ) + for name, component in pipe.components.items(): + if isinstance(component, torch.nn.Module): + assert not hasattr(component.config, "quantization_config") + + @parameterized.expand(["quant_kwargs", "quant_mapping"]) + def test_quant_config_repr(self, method): + component_name = "transformer" + if method == "quant_kwargs": + components_to_quantize = [component_name] + quant_config = PipelineQuantizationConfig( + quant_backend="bitsandbytes_8bit", + quant_kwargs={"load_in_8bit": True}, + components_to_quantize=components_to_quantize, + ) + else: + quant_config = PipelineQuantizationConfig( + quant_mapping={component_name: BitsAndBytesConfig(load_in_8bit=True)} + ) + + pipe = DiffusionPipeline.from_pretrained( + self.model_name, + quantization_config=quant_config, + torch_dtype=torch.bfloat16, + ) + assert getattr(pipe, "quantization_config", None) is not None + retrieved_config = pipe.quantization_config + expected_config = """ +transformer BitsAndBytesConfig { + "_load_in_4bit": false, + "_load_in_8bit": true, + "bnb_4bit_compute_dtype": "float32", + "bnb_4bit_quant_storage": "uint8", + "bnb_4bit_quant_type": "fp4", + "bnb_4bit_use_double_quant": false, + "llm_int8_enable_fp32_cpu_offload": false, + "llm_int8_has_fp16_weight": false, + "llm_int8_skip_modules": null, + "llm_int8_threshold": 6.0, + "load_in_4bit": false, + "load_in_8bit": true, + "quant_method": "bitsandbytes" +} + +""" + expected_data = self._parse_config_string(expected_config) + actual_data = self._parse_config_string(str(retrieved_config)) + assert actual_data == expected_data + + def _parse_config_string(self, config_string: str) -> tuple[str, dict]: + first_brace = config_string.find("{") + if first_brace == -1: + raise ValueError("Could not find opening brace '{' in the string.") + + json_part = config_string[first_brace:] + data = json.loads(json_part) + + return data + + def test_single_component_to_quantize(self): + component_to_quantize = "transformer" + quant_config = PipelineQuantizationConfig( + quant_backend="bitsandbytes_8bit", + quant_kwargs={"load_in_8bit": True}, + components_to_quantize=component_to_quantize, + ) + pipe = DiffusionPipeline.from_pretrained( + self.model_name, + quantization_config=quant_config, + torch_dtype=torch.bfloat16, + ) + for name, component in pipe.components.items(): + if name == component_to_quantize: + assert hasattr(component.config, "quantization_config") + + +# ======================== BitsAndBytes ======================== + + +# Model-level BitsAndBytes tests live in `tests/models/testing_utils/quantization.py` +# (`BitsAndBytesTesterMixin` / `BitsAndBytesCompileTesterMixin`), wired into model test files via +# concrete classes (e.g. `TestFluxTransformerBitsAndBytes`). Only pipeline-level coverage remains here. +@is_quantization +@is_bitsandbytes +@require_bitsandbytes_version_greater("0.43.2") +@require_accelerate +@require_torch +@require_torch_accelerator +@slow +class Base4bitTests: + # We need to test on relatively large models (aka >1b parameters otherwise the quantiztion may not work as expected) + # Therefore here we use only SD3 to test our module + model_name = "stabilityai/stable-diffusion-3-medium-diffusers" + + prompt = "a beautiful sunset amidst the mountains." + num_inference_steps = 10 + seed = 0 + + @pytest.fixture(autouse=True, scope="class") + def _toggle_determinism(self): + was_enabled = torch.are_deterministic_algorithms_enabled() + if not was_enabled: + torch.use_deterministic_algorithms(True) + yield + if not was_enabled: + torch.use_deterministic_algorithms(False) + + +@require_transformers_version_greater("4.44.0") +class TestSlowBnb4Bit(Base4bitTests): + @pytest.fixture(autouse=True) + def _setup_slow(self): + gc.collect() + backend_empty_cache(torch_device) + + nf4_config = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_quant_type="nf4", + bnb_4bit_compute_dtype=torch.float16, + ) + model_4bit = SD3Transformer2DModel.from_pretrained( + self.model_name, subfolder="transformer", quantization_config=nf4_config, device_map=torch_device + ) + self.pipeline_4bit = DiffusionPipeline.from_pretrained( + self.model_name, transformer=model_4bit, torch_dtype=torch.float16 + ) + self.pipeline_4bit.enable_model_cpu_offload() + yield + del self.pipeline_4bit + + gc.collect() + backend_empty_cache(torch_device) + + def test_quality(self): + output = self.pipeline_4bit( + prompt=self.prompt, + num_inference_steps=self.num_inference_steps, + generator=torch.manual_seed(self.seed), + output_type="np", + ).images + + out_slice = output[0, -3:, -3:, -1].flatten() + expected_slice = np.array([0.1123, 0.1296, 0.1609, 0.1042, 0.1230, 0.1274, 0.0928, 0.1165, 0.1216]) + + max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice) + assert max_diff < 1e-2 + + def test_generate_quality_dequantize(self): + r""" + Test that loading the model and unquantize it produce correct results. + """ + self.pipeline_4bit.transformer.dequantize() + output = self.pipeline_4bit( + prompt=self.prompt, + num_inference_steps=self.num_inference_steps, + generator=torch.manual_seed(self.seed), + output_type="np", + ).images + + out_slice = output[0, -3:, -3:, -1].flatten() + expected_slice = np.array([0.1216, 0.1387, 0.1584, 0.1152, 0.1318, 0.1282, 0.1062, 0.1226, 0.1228]) + max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice) + assert max_diff < 1e-3 + + # Since we offloaded the `pipeline_4bit.transformer` to CPU (result of `enable_model_cpu_offload()), check + # the following. + assert self.pipeline_4bit.transformer.device.type == "cpu" + # calling it again shouldn't be a problem + _ = self.pipeline_4bit( + prompt=self.prompt, + num_inference_steps=2, + generator=torch.manual_seed(self.seed), + output_type="np", + ).images + + def test_moving_to_cpu_throws_warning(self): + nf4_config = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_quant_type="nf4", + bnb_4bit_compute_dtype=torch.float16, + ) + model_4bit = SD3Transformer2DModel.from_pretrained( + self.model_name, subfolder="transformer", quantization_config=nf4_config, device_map=torch_device + ) + + logger = logging.get_logger("diffusers.pipelines.pipeline_utils") + logger.setLevel(30) + with CaptureLogger(logger) as cap_logger: + # Because `model.dtype` will return torch.float16 as SD3 transformer has + # a conv layer as the first layer. + _ = DiffusionPipeline.from_pretrained( + self.model_name, transformer=model_4bit, torch_dtype=torch.float16 + ).to("cpu") + + assert "Pipelines loaded with `dtype=torch.float16`" in cap_logger.out + + @pytest.mark.xfail( + condition=is_accelerate_version("<=", "1.1.1"), + reason="Test will pass after https://github.com/huggingface/accelerate/pull/3223 is in a release.", + strict=True, + ) + def test_pipeline_cuda_placement_works_with_nf4(self): + transformer_nf4_config = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_quant_type="nf4", + bnb_4bit_compute_dtype=torch.float16, + ) + transformer_4bit = SD3Transformer2DModel.from_pretrained( + self.model_name, + subfolder="transformer", + quantization_config=transformer_nf4_config, + torch_dtype=torch.float16, + device_map=torch_device, + ) + text_encoder_3_nf4_config = TranBitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_quant_type="nf4", + bnb_4bit_compute_dtype=torch.float16, + ) + text_encoder_3_4bit = T5EncoderModel.from_pretrained( + self.model_name, + subfolder="text_encoder_3", + quantization_config=text_encoder_3_nf4_config, + torch_dtype=torch.float16, + device_map=torch_device, + ) + # CUDA device placement works. + pipeline_4bit = DiffusionPipeline.from_pretrained( + self.model_name, + transformer=transformer_4bit, + text_encoder_3=text_encoder_3_4bit, + torch_dtype=torch.float16, + ).to(torch_device) + + # Check if inference works. + _ = pipeline_4bit(self.prompt, max_sequence_length=20, num_inference_steps=2) + + del pipeline_4bit + + +@require_transformers_version_greater("4.44.0") +class TestSlowBnb4BitFlux(Base4bitTests): + @pytest.fixture(autouse=True) + def _setup_flux(self): + gc.collect() + backend_empty_cache(torch_device) + + model_id = "hf-internal-testing/flux.1-dev-nf4-pkg" + t5_4bit = T5EncoderModel.from_pretrained(model_id, subfolder="text_encoder_2") + transformer_4bit = FluxTransformer2DModel.from_pretrained(model_id, subfolder="transformer") + self.pipeline_4bit = DiffusionPipeline.from_pretrained( + "black-forest-labs/FLUX.1-dev", + text_encoder_2=t5_4bit, + transformer=transformer_4bit, + torch_dtype=torch.float16, + ) + self.pipeline_4bit.enable_model_cpu_offload() + yield + del self.pipeline_4bit + + gc.collect() + backend_empty_cache(torch_device) + + def test_quality(self): + # keep the resolution and max tokens to a lower number for faster execution. + output = self.pipeline_4bit( + prompt=self.prompt, + num_inference_steps=self.num_inference_steps, + generator=torch.manual_seed(self.seed), + height=256, + width=256, + max_sequence_length=64, + output_type="np", + ).images + + out_slice = output[0, -3:, -3:, -1].flatten() + expected_slice = np.array([0.0583, 0.0586, 0.0632, 0.0815, 0.0813, 0.0947, 0.1040, 0.1145, 0.1265]) + + max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice) + assert max_diff < 1e-3 + + @require_peft_backend + def test_lora_loading(self): + self.pipeline_4bit.load_lora_weights( + hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors"), adapter_name="hyper-sd" + ) + self.pipeline_4bit.set_adapters("hyper-sd", adapter_weights=0.125) + + output = self.pipeline_4bit( + prompt=self.prompt, + height=256, + width=256, + max_sequence_length=64, + output_type="np", + num_inference_steps=8, + generator=torch.Generator().manual_seed(42), + ).images + out_slice = output[0, -3:, -3:, -1].flatten() + expected_slice = np.array([0.5347, 0.5342, 0.5283, 0.5093, 0.4988, 0.5093, 0.5044, 0.5015, 0.4946]) + + max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice) + assert max_diff < 1e-3 + + +@require_transformers_version_greater("4.44.0") +@require_peft_backend +class TestSlowBnb4BitFluxControlWithLora(Base4bitTests): + @pytest.fixture(autouse=True) + def _setup_flux_control(self): + gc.collect() + backend_empty_cache(torch_device) + + self.pipeline_4bit = FluxControlPipeline.from_pretrained("eramth/flux-4bit", torch_dtype=torch.float16) + self.pipeline_4bit.enable_model_cpu_offload() + yield + del self.pipeline_4bit + + gc.collect() + backend_empty_cache(torch_device) + + def test_lora_loading(self): + self.pipeline_4bit.load_lora_weights("black-forest-labs/FLUX.1-Canny-dev-lora") + + output = self.pipeline_4bit( + prompt=self.prompt, + control_image=Image.new(mode="RGB", size=(256, 256)), + height=256, + width=256, + max_sequence_length=64, + output_type="np", + num_inference_steps=8, + generator=torch.Generator().manual_seed(42), + ).images + out_slice = output[0, -3:, -3:, -1].flatten() + expected_slice = np.array([0.1636, 0.1675, 0.1982, 0.1743, 0.1809, 0.1936, 0.1743, 0.2095, 0.2139]) + + max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice) + assert max_diff < 1e-3, f"{out_slice=} != {expected_slice=}" + + +@is_quantization +@is_bitsandbytes +@require_torch_version_greater("2.7.1") +@require_bitsandbytes_version_greater("0.45.5") +class TestBnb4BitCompile(QuantCompileTests): + @property + def quantization_config(self): + return PipelineQuantizationConfig( + quant_backend="bitsandbytes_4bit", + quant_kwargs={ + "load_in_4bit": True, + "bnb_4bit_quant_type": "nf4", + "bnb_4bit_compute_dtype": torch.bfloat16, + }, + components_to_quantize=["transformer", "text_encoder_2"], + ) + + +@is_quantization +@is_bitsandbytes +@require_bitsandbytes_version_greater("0.43.2") +@require_accelerate +@require_torch +@require_torch_accelerator +@slow +class Base8bitTests: + # We need to test on relatively large models (aka >1b parameters otherwise the quantiztion may not work as expected) + # Therefore here we use only SD3 to test our module + model_name = "stabilityai/stable-diffusion-3-medium-diffusers" + + prompt = "a beautiful sunset amidst the mountains." + num_inference_steps = 10 + seed = 0 + + @pytest.fixture(autouse=True, scope="class") + def _toggle_determinism(self): + was_enabled = torch.are_deterministic_algorithms_enabled() + if not was_enabled: + torch.use_deterministic_algorithms(True) + yield + if not was_enabled: + torch.use_deterministic_algorithms(False) + + +@require_transformers_version_greater("4.44.0") +class TestSlowBnb8bit(Base8bitTests): + @pytest.fixture(autouse=True) + def _setup_slow(self): + gc.collect() + backend_empty_cache(torch_device) + + mixed_int8_config = BitsAndBytesConfig(load_in_8bit=True) + model_8bit = SD3Transformer2DModel.from_pretrained( + self.model_name, subfolder="transformer", quantization_config=mixed_int8_config, device_map=torch_device + ) + self.pipeline_8bit = DiffusionPipeline.from_pretrained( + self.model_name, transformer=model_8bit, torch_dtype=torch.float16 + ) + self.pipeline_8bit.enable_model_cpu_offload() + yield + del self.pipeline_8bit + + gc.collect() + backend_empty_cache(torch_device) + + def test_quality(self): + output = self.pipeline_8bit( + prompt=self.prompt, + num_inference_steps=self.num_inference_steps, + generator=torch.manual_seed(self.seed), + output_type="np", + ).images + out_slice = output[0, -3:, -3:, -1].flatten() + expected_slice = np.array([0.0674, 0.0623, 0.0364, 0.0632, 0.0671, 0.0430, 0.0317, 0.0493, 0.0583]) + + max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice) + assert max_diff < 1e-2 + + def test_model_cpu_offload_raises_warning(self): + model_8bit = SD3Transformer2DModel.from_pretrained( + self.model_name, + subfolder="transformer", + quantization_config=BitsAndBytesConfig(load_in_8bit=True), + device_map=torch_device, + ) + pipeline_8bit = DiffusionPipeline.from_pretrained( + self.model_name, transformer=model_8bit, torch_dtype=torch.float16 + ) + logger = logging.get_logger("diffusers.pipelines.pipeline_utils") + logger.setLevel(30) + + with CaptureLogger(logger) as cap_logger: + pipeline_8bit.enable_model_cpu_offload() + + assert "has been loaded in `bitsandbytes` 8bit" in cap_logger.out + + def test_moving_to_cpu_throws_warning(self): + model_8bit = SD3Transformer2DModel.from_pretrained( + self.model_name, + subfolder="transformer", + quantization_config=BitsAndBytesConfig(load_in_8bit=True), + device_map=torch_device, + ) + logger = logging.get_logger("diffusers.pipelines.pipeline_utils") + logger.setLevel(30) + + with CaptureLogger(logger) as cap_logger: + # Because `model.dtype` will return torch.float16 as SD3 transformer has + # a conv layer as the first layer. + _ = DiffusionPipeline.from_pretrained( + self.model_name, transformer=model_8bit, torch_dtype=torch.float16 + ).to("cpu") + + assert "Pipelines loaded with `dtype=torch.float16`" in cap_logger.out + + def test_generate_quality_dequantize(self): + r""" + Test that loading the model and unquantize it produce correct results. + """ + self.pipeline_8bit.transformer.dequantize() + output = self.pipeline_8bit( + prompt=self.prompt, + num_inference_steps=self.num_inference_steps, + generator=torch.manual_seed(self.seed), + output_type="np", + ).images + + out_slice = output[0, -3:, -3:, -1].flatten() + expected_slice = np.array([0.0266, 0.0264, 0.0271, 0.0110, 0.0310, 0.0098, 0.0078, 0.0256, 0.0208]) + max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice) + assert max_diff < 1e-2 + + # 8bit models cannot be offloaded to CPU. + assert self.pipeline_8bit.transformer.device.type == torch_device + # calling it again shouldn't be a problem + _ = self.pipeline_8bit( + prompt=self.prompt, + num_inference_steps=2, + generator=torch.manual_seed(self.seed), + output_type="np", + ).images + + @pytest.mark.xfail( + condition=is_accelerate_version("<=", "1.1.1"), + reason="Test will pass after https://github.com/huggingface/accelerate/pull/3223 is in a release.", + strict=True, + ) + def test_pipeline_cuda_placement_works_with_mixed_int8(self): + transformer_8bit_config = BitsAndBytesConfig(load_in_8bit=True) + transformer_8bit = SD3Transformer2DModel.from_pretrained( + self.model_name, + subfolder="transformer", + quantization_config=transformer_8bit_config, + torch_dtype=torch.float16, + device_map=torch_device, + ) + text_encoder_3_8bit_config = TranBitsAndBytesConfig(load_in_8bit=True) + text_encoder_3_8bit = T5EncoderModel.from_pretrained( + self.model_name, + subfolder="text_encoder_3", + quantization_config=text_encoder_3_8bit_config, + torch_dtype=torch.float16, + device_map=torch_device, + ) + + # CUDA device placement works. + device = torch_device if torch_device != "rocm" else "cuda" + pipeline_8bit = DiffusionPipeline.from_pretrained( + self.model_name, + transformer=transformer_8bit, + text_encoder_3=text_encoder_3_8bit, + torch_dtype=torch.float16, + ).to(device) + + # Check if inference works. + _ = pipeline_8bit(self.prompt, max_sequence_length=20, num_inference_steps=2) + + del pipeline_8bit + + +@require_transformers_version_greater("4.44.0") +@require_big_accelerator +class TestSlowBnb8bitFlux(Base8bitTests): + @pytest.fixture(autouse=True) + def _setup_slow_flux(self): + gc.collect() + backend_empty_cache(torch_device) + + model_id = "hf-internal-testing/flux.1-dev-int8-pkg" + t5_8bit = T5EncoderModel.from_pretrained(model_id, subfolder="text_encoder_2") + transformer_8bit = FluxTransformer2DModel.from_pretrained(model_id, subfolder="transformer") + self.pipeline_8bit = DiffusionPipeline.from_pretrained( + "black-forest-labs/FLUX.1-dev", + text_encoder_2=t5_8bit, + transformer=transformer_8bit, + torch_dtype=torch.float16, + ) + self.pipeline_8bit.enable_model_cpu_offload() + yield + del self.pipeline_8bit + + gc.collect() + backend_empty_cache(torch_device) + + def test_quality(self): + # keep the resolution and max tokens to a lower number for faster execution. + output = self.pipeline_8bit( + prompt=self.prompt, + num_inference_steps=self.num_inference_steps, + generator=torch.manual_seed(self.seed), + height=256, + width=256, + max_sequence_length=64, + output_type="np", + ).images + out_slice = output[0, -3:, -3:, -1].flatten() + expected_slice = np.array([0.0574, 0.0554, 0.0581, 0.0686, 0.0676, 0.0759, 0.0757, 0.0803, 0.0930]) + + max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice) + assert max_diff < 1e-3 + + @require_peft_version_greater("0.14.0") + def test_lora_loading(self): + self.pipeline_8bit.load_lora_weights( + hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors"), adapter_name="hyper-sd" + ) + self.pipeline_8bit.set_adapters("hyper-sd", adapter_weights=0.125) + + output = self.pipeline_8bit( + prompt=self.prompt, + height=256, + width=256, + max_sequence_length=64, + output_type="np", + num_inference_steps=8, + generator=torch.manual_seed(42), + ).images + out_slice = output[0, -3:, -3:, -1].flatten() + + expected_slice = np.array([0.3916, 0.3916, 0.3887, 0.4243, 0.4155, 0.4233, 0.4570, 0.4531, 0.4248]) + + max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice) + assert max_diff < 1e-3 + + +@require_transformers_version_greater("4.44.0") +@require_peft_backend +class TestSlowBnb8bitFluxControlWithLora(Base8bitTests): + @pytest.fixture(autouse=True) + def _setup_flux_control_lora(self): + gc.collect() + backend_empty_cache(torch_device) + + self.pipeline_8bit = FluxControlPipeline.from_pretrained( + "black-forest-labs/FLUX.1-dev", + quantization_config=PipelineQuantizationConfig( + quant_backend="bitsandbytes_8bit", + quant_kwargs={"load_in_8bit": True}, + components_to_quantize=["transformer", "text_encoder_2"], + ), + torch_dtype=torch.float16, + ) + self.pipeline_8bit.enable_model_cpu_offload() + yield + del self.pipeline_8bit + + gc.collect() + backend_empty_cache(torch_device) + + def test_lora_loading(self): + self.pipeline_8bit.load_lora_weights("black-forest-labs/FLUX.1-Canny-dev-lora") + + output = self.pipeline_8bit( + prompt=self.prompt, + control_image=Image.new(mode="RGB", size=(256, 256)), + height=256, + width=256, + max_sequence_length=64, + output_type="np", + num_inference_steps=8, + generator=torch.Generator().manual_seed(42), + ).images + out_slice = output[0, -3:, -3:, -1].flatten() + # Hardware-dependent: the Control LoRA dequantizes and expands `x_embedder`, and the error + # accumulates over the 8 denoising steps enough that even different CUDA GPUs disagree, so + # reference slices are stored per accelerator backend. + expected_slices = Expectations( + { + (None, None): np.array([0.2029, 0.2136, 0.2268, 0.1921, 0.1997, 0.2185, 0.2021, 0.2183, 0.2292]), + ("xpu", 5): np.array([0.0955, 0.1223, 0.1509, 0.0872, 0.1155, 0.1890, 0.0754, 0.1028, 0.2178]), + } + ) + expected_slice = expected_slices.get_expectation() + + max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice) + assert max_diff < 1e-3, f"{out_slice=} != {expected_slice=}" + + +@is_quantization +@is_bitsandbytes +@require_torch_version_greater_equal("2.6.0") +@require_bitsandbytes_version_greater("0.48.0") +class TestBnb8BitCompile(QuantCompileTests): + @property + def quantization_config(self): + return PipelineQuantizationConfig( + quant_backend="bitsandbytes_8bit", + quant_kwargs={"load_in_8bit": True}, + components_to_quantize=["transformer", "text_encoder_2"], + ) + + def test_torch_compile_with_cpu_offload(self): + super()._test_torch_compile_with_cpu_offload(torch_dtype=torch.float16) + + +# ======================== Quanto ======================== + + +# ======================== TorchAO ======================== + + +def _is_xpu_or_cuda_capability_atleast_8_9() -> bool: + if torch.cuda.is_available(): + major, minor = torch.cuda.get_device_capability() + if major == 8: + return minor >= 9 + return major >= 9 + elif torch.xpu.is_available(): + return True + return False + + +# Model-level TorchAO tests live in `tests/models/testing_utils/quantization.py` +# (`TorchAoTesterMixin` / `TorchAoCompileTesterMixin`), wired into model test files via concrete +# classes (e.g. `TestFluxTransformerTorchAo`). Only pipeline-level coverage remains here. +# Slices for these tests have been obtained on our aws-g6e-xlarge-plus runners +@is_quantization +@is_torchao +@require_torch +@require_torch_accelerator +@require_torchao_version_greater_or_equal("0.15.0") +class TestTorchAo: + @pytest.fixture(autouse=True) + def _setup_torchao(self): + yield + gc.collect() + backend_empty_cache(torch_device) + + def get_dummy_components( + self, quantization_config: TorchAoConfig, model_id: str = "hf-internal-testing/tiny-flux-pipe" + ): + transformer = FluxTransformer2DModel.from_pretrained( + model_id, + subfolder="transformer", + quantization_config=quantization_config, + torch_dtype=torch.bfloat16, + ) + text_encoder = CLIPTextModel.from_pretrained(model_id, subfolder="text_encoder", torch_dtype=torch.bfloat16) + text_encoder_2 = T5EncoderModel.from_pretrained( + model_id, subfolder="text_encoder_2", torch_dtype=torch.bfloat16 + ) + tokenizer = CLIPTokenizer.from_pretrained(model_id, subfolder="tokenizer") + tokenizer_2 = AutoTokenizer.from_pretrained(model_id, subfolder="tokenizer_2") + vae = AutoencoderKL.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.bfloat16) + scheduler = FlowMatchEulerDiscreteScheduler() + + return { + "scheduler": scheduler, + "text_encoder": text_encoder, + "text_encoder_2": text_encoder_2, + "tokenizer": tokenizer, + "tokenizer_2": tokenizer_2, + "transformer": transformer, + "vae": vae, + } + + def get_dummy_inputs(self, device: torch.device, seed: int = 0): + if str(device).startswith("mps"): + generator = torch.manual_seed(seed) + else: + generator = torch.Generator().manual_seed(seed) + + inputs = { + "prompt": "an astronaut riding a horse in space", + "height": 32, + "width": 32, + "num_inference_steps": 2, + "output_type": "np", + "generator": generator, + } + + return inputs + + def _test_quant_type(self, quantization_config: TorchAoConfig, expected_slice: list[float], model_id: str): + components = self.get_dummy_components(quantization_config, model_id) + pipe = FluxPipeline(**components) + pipe.to(device=torch_device) + + inputs = self.get_dummy_inputs(torch_device) + output = pipe(**inputs)[0] + output_slice = output[-1, -1, -3:, -3:].flatten() + + assert np.allclose(output_slice, expected_slice, atol=1e-3, rtol=1e-3) + + def test_quantization(self): + for model_id in ["hf-internal-testing/tiny-flux-pipe", "hf-internal-testing/tiny-flux-sharded"]: + # fmt: off + QUANTIZATION_TYPES_TO_TEST = [ + (Int4WeightOnlyConfig(version=2), np.array([0.4648, 0.5234, 0.5547, 0.4219, 0.4414, 0.6445, 0.4336, 0.4531, 0.5625])), + (Int8DynamicActivationIntxWeightConfig(version=2), np.array([0.4688, 0.5195, 0.5547, 0.418, 0.4414, 0.6406, 0.4336, 0.4531, 0.5625])), + (Int8WeightOnlyConfig(version=2), np.array([0.4648, 0.5195, 0.5547, 0.4199, 0.4414, 0.6445, 0.4316, 0.4531, 0.5625])), + (Int8DynamicActivationInt8WeightConfig(version=2), np.array([0.4648, 0.5195, 0.5547, 0.4199, 0.4414, 0.6445, 0.4316, 0.4531, 0.5625])), + (IntxWeightOnlyConfig(dtype=torch.uint4, group_size=16, version=2), np.array([0.4609, 0.5234, 0.5508, 0.4199, 0.4336, 0.6406, 0.4316, 0.4531, 0.5625])), + (IntxWeightOnlyConfig(dtype=torch.uint7, group_size=16, version=2), np.array([0.4648, 0.5195, 0.5547, 0.4219, 0.4414, 0.6445, 0.4316, 0.4531, 0.5625])), + ] + + if _is_xpu_or_cuda_capability_atleast_8_9(): + QUANTIZATION_TYPES_TO_TEST.extend([ + (Float8WeightOnlyConfig(weight_dtype=torch.float8_e5m2), np.array([0.4590, 0.5273, 0.5547, 0.4219, 0.4375, 0.6406, 0.4316, 0.4512, 0.5625])), + (Float8WeightOnlyConfig(weight_dtype=torch.float8_e4m3fn), np.array([0.4648, 0.5234, 0.5547, 0.4219, 0.4414, 0.6406, 0.4316, 0.4531, 0.5625])), + ]) + # fmt: on + + for quant_config, expected_slice in QUANTIZATION_TYPES_TO_TEST: + quantization_config = TorchAoConfig(quant_type=quant_config, modules_to_not_convert=["x_embedder"]) + self._test_quant_type(quantization_config, expected_slice, model_id) + + def test_sequential_cpu_offload(self): + r""" + A test that checks if inference runs as expected when sequential cpu offloading is enabled. + """ + quantization_config = TorchAoConfig(Int8WeightOnlyConfig()) + components = self.get_dummy_components(quantization_config) + pipe = FluxPipeline(**components) + pipe.enable_sequential_cpu_offload() + + inputs = self.get_dummy_inputs(torch_device) + _ = pipe(**inputs) + + @require_torchao_version_greater_or_equal("0.15.0") + def test_aobase_config(self): + quantization_config = TorchAoConfig(Int8WeightOnlyConfig()) + components = self.get_dummy_components(quantization_config) + pipe = FluxPipeline(**components).to(torch_device) + + inputs = self.get_dummy_inputs(torch_device) + _ = pipe(**inputs) + + +@is_quantization +@is_torchao +@require_torchao_version_greater_or_equal("0.15.0") +class TestTorchAoCompile(QuantCompileTests): + @property + def quantization_config(self): + return PipelineQuantizationConfig( + quant_mapping={"transformer": TorchAoConfig(Int8WeightOnlyConfig())}, + ) + + def test_torch_compile_with_cpu_offload(self): + pipe = self._init_pipeline(self.quantization_config, torch.bfloat16) + pipe.enable_model_cpu_offload() + # No compilation because it fails with: + # RuntimeError: _apply(): Couldn't swap Linear.weight + + # small resolutions to ensure speedy execution. + pipe("a dog", num_inference_steps=2, max_sequence_length=16, height=256, width=256) + + +# Slices for these tests have been obtained on our aws-g6e-xlarge-plus runners +@is_quantization +@is_torchao +@require_torch +@require_torch_accelerator +@require_torchao_version_greater_or_equal("0.15.0") +@slow +@nightly +class TestSlowTorchAo: + @pytest.fixture(autouse=True) + def _setup_slow_torchao(self): + yield + gc.collect() + backend_empty_cache(torch_device) + + def get_dummy_components(self, quantization_config: TorchAoConfig): + # This is just for convenience, so that we can modify it at one place for custom environments and locally testing + cache_dir = None + model_id = "black-forest-labs/FLUX.1-dev" + transformer = FluxTransformer2DModel.from_pretrained( + model_id, + subfolder="transformer", + quantization_config=quantization_config, + torch_dtype=torch.bfloat16, + cache_dir=cache_dir, + ) + text_encoder = CLIPTextModel.from_pretrained( + model_id, subfolder="text_encoder", torch_dtype=torch.bfloat16, cache_dir=cache_dir + ) + text_encoder_2 = T5EncoderModel.from_pretrained( + model_id, subfolder="text_encoder_2", torch_dtype=torch.bfloat16, cache_dir=cache_dir + ) + tokenizer = CLIPTokenizer.from_pretrained(model_id, subfolder="tokenizer", cache_dir=cache_dir) + tokenizer_2 = AutoTokenizer.from_pretrained(model_id, subfolder="tokenizer_2", cache_dir=cache_dir) + vae = AutoencoderKL.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.bfloat16, cache_dir=cache_dir) + scheduler = FlowMatchEulerDiscreteScheduler() + + return { + "scheduler": scheduler, + "text_encoder": text_encoder, + "text_encoder_2": text_encoder_2, + "tokenizer": tokenizer, + "tokenizer_2": tokenizer_2, + "transformer": transformer, + "vae": vae, + } + + def get_dummy_inputs(self, device: torch.device, seed: int = 0): + if str(device).startswith("mps"): + generator = torch.manual_seed(seed) + else: + generator = torch.Generator().manual_seed(seed) + + inputs = { + "prompt": "an astronaut riding a horse in space", + "height": 512, + "width": 512, + "num_inference_steps": 20, + "output_type": "np", + "generator": generator, + } + + return inputs + + def _test_quant_type(self, quantization_config, expected_slice): + components = self.get_dummy_components(quantization_config) + pipe = FluxPipeline(**components) + pipe.enable_model_cpu_offload() + + weight = pipe.transformer.transformer_blocks[0].ff.net[2].weight + assert isinstance(weight, TorchAOBaseTensor) + + inputs = self.get_dummy_inputs(torch_device) + output = pipe(**inputs)[0].flatten() + output_slice = np.concatenate((output[:16], output[-16:])) + assert np.allclose(output_slice, expected_slice, atol=1e-3, rtol=1e-3) + + def test_quantization(self): + # fmt: off + QUANTIZATION_TYPES_TO_TEST = [ + (Int8WeightOnlyConfig(), np.array([0.0505, 0.0742, 0.1367, 0.0429, 0.0585, 0.1386, 0.0585, 0.0703, 0.1367, 0.0566, 0.0703, 0.1464, 0.0546, 0.0703, 0.1425, 0.0546, 0.3535, 0.7578, 0.5000, 0.4062, 0.7656, 0.5117, 0.4121, 0.7656, 0.5117, 0.3984, 0.7578, 0.5234, 0.4023, 0.7382, 0.5390, 0.4570])), + (Int8DynamicActivationInt8WeightConfig(), np.array([0.0546, 0.0761, 0.1386, 0.0488, 0.0644, 0.1425, 0.0605, 0.0742, 0.1406, 0.0625, 0.0722, 0.1523, 0.0625, 0.0742, 0.1503, 0.0605, 0.3886, 0.7968, 0.5507, 0.4492, 0.7890, 0.5351, 0.4316, 0.8007, 0.5390, 0.4179, 0.8281, 0.5820, 0.4531, 0.7812, 0.5703, 0.4921])), + ] + + if _is_xpu_or_cuda_capability_atleast_8_9(): + QUANTIZATION_TYPES_TO_TEST.extend([ + (Float8WeightOnlyConfig(weight_dtype=torch.float8_e4m3fn), np.array([0.0546, 0.0722, 0.1328, 0.0468, 0.0585, 0.1367, 0.0605, 0.0703, 0.1328, 0.0625, 0.0703, 0.1445, 0.0585, 0.0703, 0.1406, 0.0605, 0.3496, 0.7109, 0.4843, 0.4042, 0.7226, 0.5000, 0.4160, 0.7031, 0.4824, 0.3886, 0.6757, 0.4667, 0.3710, 0.6679, 0.4902, 0.4238])), + ]) + # fmt: on + + for quant_config, expected_slice in QUANTIZATION_TYPES_TO_TEST: + quantization_config = TorchAoConfig(quant_type=quant_config, modules_to_not_convert=["x_embedder"]) + self._test_quant_type(quantization_config, expected_slice) + gc.collect() + backend_empty_cache(torch_device) + backend_synchronize(torch_device) + + def test_serialization_int8wo(self): + quantization_config = TorchAoConfig(Int8WeightOnlyConfig()) + components = self.get_dummy_components(quantization_config) + pipe = FluxPipeline(**components) + pipe.enable_model_cpu_offload() + + weight = pipe.transformer.x_embedder.weight + assert isinstance(weight, Int8Tensor) + + inputs = self.get_dummy_inputs(torch_device) + output = pipe(**inputs)[0].flatten()[:128] + + with tempfile.TemporaryDirectory() as tmp_dir: + pipe.transformer.save_pretrained(tmp_dir, safe_serialization=False) + pipe.remove_all_hooks() + del pipe.transformer + gc.collect() + backend_empty_cache(torch_device) + backend_synchronize(torch_device) + transformer = FluxTransformer2DModel.from_pretrained( + tmp_dir, torch_dtype=torch.bfloat16, use_safetensors=False + ) + pipe.transformer = transformer + pipe.enable_model_cpu_offload() + + weight = transformer.x_embedder.weight + assert isinstance(weight, Int8Tensor) + + loaded_output = pipe(**inputs)[0].flatten()[:128] + # Seems to require higher tolerance depending on which machine it is being run. + # A difference of 0.06 in normalized pixel space (-1 to 1), corresponds to a difference of + # 0.06 / 2 * 255 = 7.65 in pixel space (0 to 255). On our CI runners, the difference is about 0.04, + # on DGX it is 0.06, and on audace it is 0.037. So, we are using a tolerance of 0.06 here. + assert np.allclose(output, loaded_output, atol=0.06) + + +# ======================== GGUF ======================== + + +# Model-level GGUF tests live in `tests/models/testing_utils/quantization.py` +# (`GGUFTesterMixin` / `GGUFCompileTesterMixin`) and backend-level ones (quantized parameter/layer +# inspection, memory use, CUDA kernels) in `tests/quantization/gguf/`. Only pipeline-level coverage +# remains here. +@is_quantization +@is_gguf +@nightly +@require_big_accelerator +@require_accelerate +@require_gguf_version_greater_or_equal("0.10.0") +class GGUFPipelineTests: + @pytest.fixture(autouse=True) + def _cleanup(self): + gc.collect() + backend_empty_cache(torch_device) + yield + gc.collect() + backend_empty_cache(torch_device) + + +class TestFluxGGUFPipeline(GGUFPipelineTests): + ckpt_path = "https://huggingface.co/city96/FLUX.1-dev-gguf/blob/main/flux1-dev-Q2_K.gguf" + model_cls = FluxTransformer2DModel + torch_dtype = torch.bfloat16 + + def test_pipeline_inference(self): + quantization_config = GGUFQuantizationConfig(compute_dtype=self.torch_dtype) + transformer = self.model_cls.from_single_file( + self.ckpt_path, quantization_config=quantization_config, torch_dtype=self.torch_dtype + ) + pipe = FluxPipeline.from_pretrained( + "black-forest-labs/FLUX.1-dev", transformer=transformer, torch_dtype=self.torch_dtype + ) + pipe.enable_model_cpu_offload() + + prompt = "a cat holding a sign that says hello" + output = pipe( + prompt=prompt, num_inference_steps=2, generator=torch.Generator("cpu").manual_seed(0), output_type="np" + ).images[0] + output_slice = output[:3, :3, :].flatten() + expected_slice = np.array( + [ + 0.47265625, + 0.43359375, + 0.359375, + 0.47070312, + 0.421875, + 0.34375, + 0.46875, + 0.421875, + 0.34765625, + 0.46484375, + 0.421875, + 0.34179688, + 0.47070312, + 0.42578125, + 0.34570312, + 0.46875, + 0.42578125, + 0.3515625, + 0.45507812, + 0.4140625, + 0.33984375, + 0.4609375, + 0.41796875, + 0.34375, + 0.45898438, + 0.41796875, + 0.34375, + ] + ) + max_diff = numpy_cosine_similarity_distance(expected_slice, output_slice) + assert max_diff < 1e-4 + + +class TestSD35LargeGGUFPipeline(GGUFPipelineTests): + ckpt_path = "https://huggingface.co/city96/stable-diffusion-3.5-large-gguf/blob/main/sd3.5_large-Q4_0.gguf" + model_cls = SD3Transformer2DModel + torch_dtype = torch.bfloat16 + + def test_pipeline_inference(self): + quantization_config = GGUFQuantizationConfig(compute_dtype=self.torch_dtype) + transformer = self.model_cls.from_single_file( + self.ckpt_path, quantization_config=quantization_config, torch_dtype=self.torch_dtype + ) + pipe = StableDiffusion3Pipeline.from_pretrained( + "stabilityai/stable-diffusion-3.5-large", transformer=transformer, torch_dtype=self.torch_dtype + ) + pipe.enable_model_cpu_offload() + + prompt = "a cat holding a sign that says hello" + output = pipe( + prompt=prompt, + num_inference_steps=2, + generator=torch.Generator("cpu").manual_seed(0), + output_type="np", + ).images[0] + output_slice = output[:3, :3, :].flatten() + expected_slices = Expectations( + { + ("xpu", 3): np.array( + [ + 0.16796875, + 0.27929688, + 0.28320312, + 0.11328125, + 0.27539062, + 0.26171875, + 0.10742188, + 0.26367188, + 0.26171875, + 0.1484375, + 0.2734375, + 0.296875, + 0.13476562, + 0.2890625, + 0.30078125, + 0.1171875, + 0.28125, + 0.28125, + 0.16015625, + 0.31445312, + 0.30078125, + 0.15625, + 0.32421875, + 0.296875, + 0.14453125, + 0.30859375, + 0.2890625, + ] + ), + ("cuda", 7): np.array( + [ + 0.17578125, + 0.27539062, + 0.27734375, + 0.11914062, + 0.26953125, + 0.25390625, + 0.109375, + 0.25390625, + 0.25, + 0.15039062, + 0.26171875, + 0.28515625, + 0.13671875, + 0.27734375, + 0.28515625, + 0.12109375, + 0.26757812, + 0.265625, + 0.16210938, + 0.29882812, + 0.28515625, + 0.15625, + 0.30664062, + 0.27734375, + 0.14648438, + 0.29296875, + 0.26953125, + ] + ), + } + ) + expected_slice = expected_slices.get_expectation() + max_diff = numpy_cosine_similarity_distance(expected_slice, output_slice) + assert max_diff < 1e-4 + + +class TestSD35MediumGGUFPipeline(GGUFPipelineTests): + ckpt_path = "https://huggingface.co/city96/stable-diffusion-3.5-medium-gguf/blob/main/sd3.5_medium-Q3_K_M.gguf" + model_cls = SD3Transformer2DModel + torch_dtype = torch.bfloat16 + + def test_pipeline_inference(self): + quantization_config = GGUFQuantizationConfig(compute_dtype=self.torch_dtype) + transformer = self.model_cls.from_single_file( + self.ckpt_path, quantization_config=quantization_config, torch_dtype=self.torch_dtype + ) + pipe = StableDiffusion3Pipeline.from_pretrained( + "stabilityai/stable-diffusion-3.5-medium", transformer=transformer, torch_dtype=self.torch_dtype + ) + pipe.enable_model_cpu_offload() + + prompt = "a cat holding a sign that says hello" + output = pipe( + prompt=prompt, num_inference_steps=2, generator=torch.Generator("cpu").manual_seed(0), output_type="np" + ).images[0] + output_slice = output[:3, :3, :].flatten() + expected_slice = np.array( + [ + 0.625, + 0.6171875, + 0.609375, + 0.65625, + 0.65234375, + 0.640625, + 0.6484375, + 0.640625, + 0.625, + 0.6484375, + 0.63671875, + 0.6484375, + 0.66796875, + 0.65625, + 0.65234375, + 0.6640625, + 0.6484375, + 0.6328125, + 0.6640625, + 0.6484375, + 0.640625, + 0.67578125, + 0.66015625, + 0.62109375, + 0.671875, + 0.65625, + 0.62109375, + ] + ) + max_diff = numpy_cosine_similarity_distance(expected_slice, output_slice) + assert max_diff < 1e-4 + + +class TestAuraFlowGGUFPipeline(GGUFPipelineTests): + ckpt_path = "https://huggingface.co/city96/AuraFlow-v0.3-gguf/blob/main/aura_flow_0.3-Q2_K.gguf" + model_cls = AuraFlowTransformer2DModel + torch_dtype = torch.bfloat16 + + def test_pipeline_inference(self): + quantization_config = GGUFQuantizationConfig(compute_dtype=self.torch_dtype) + transformer = self.model_cls.from_single_file( + self.ckpt_path, quantization_config=quantization_config, torch_dtype=self.torch_dtype + ) + pipe = AuraFlowPipeline.from_pretrained( + "fal/AuraFlow-v0.3", transformer=transformer, torch_dtype=self.torch_dtype + ) + pipe.enable_model_cpu_offload() + + prompt = "a pony holding a sign that says hello" + output = pipe( + prompt=prompt, num_inference_steps=2, generator=torch.Generator("cpu").manual_seed(0), output_type="np" + ).images[0] + output_slice = output[:3, :3, :].flatten() + expected_slice = np.array( + [ + 0.46484375, + 0.546875, + 0.64453125, + 0.48242188, + 0.53515625, + 0.59765625, + 0.47070312, + 0.5078125, + 0.5703125, + 0.42773438, + 0.50390625, + 0.5703125, + 0.47070312, + 0.515625, + 0.57421875, + 0.45898438, + 0.48632812, + 0.53515625, + 0.4453125, + 0.5078125, + 0.56640625, + 0.47851562, + 0.5234375, + 0.57421875, + 0.48632812, + 0.5234375, + 0.56640625, + ] + ) + max_diff = numpy_cosine_similarity_distance(expected_slice, output_slice) + assert max_diff < 1e-4 + + +@is_quantization +@is_gguf +@require_peft_backend +@nightly +@require_big_accelerator +@require_accelerate +@require_gguf_version_greater_or_equal("0.10.0") +class TestFluxControlLoRAGGUF: + def test_lora_loading(self): + ckpt_path = "https://huggingface.co/city96/FLUX.1-dev-gguf/blob/main/flux1-dev-Q2_K.gguf" + transformer = FluxTransformer2DModel.from_single_file( + ckpt_path, + quantization_config=GGUFQuantizationConfig(compute_dtype=torch.bfloat16), + torch_dtype=torch.bfloat16, + ) + pipe = FluxControlPipeline.from_pretrained( + "black-forest-labs/FLUX.1-dev", + transformer=transformer, + torch_dtype=torch.bfloat16, + ).to(torch_device) + pipe.load_lora_weights("black-forest-labs/FLUX.1-Canny-dev-lora") + + prompt = "A robot made of exotic candies and chocolates of different kinds. The background is filled with confetti and celebratory gifts." + control_image = load_image( + "https://huggingface.co/datasets/sayakpaul/sample-datasets/resolve/main/control_image_robot_canny.png" + ) + + output = pipe( + prompt=prompt, + control_image=control_image, + height=256, + width=256, + num_inference_steps=10, + guidance_scale=30.0, + output_type="np", + generator=torch.manual_seed(0), + ).images + + out_slice = output[0, -3:, -3:, -1].flatten() + expected_slice = np.array([0.8047, 0.8359, 0.8711, 0.6875, 0.7070, 0.7383, 0.5469, 0.5820, 0.6641]) + + max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice) + assert max_diff < 1e-3 + + +@is_quantization +@is_gguf +@require_torch_version_greater("2.7.1") +class TestGGUFCompile(QuantCompileTests): + torch_dtype = torch.bfloat16 + gguf_ckpt = "https://huggingface.co/city96/FLUX.1-dev-gguf/blob/main/flux1-dev-Q2_K.gguf" + + @property + def quantization_config(self): + return GGUFQuantizationConfig(compute_dtype=self.torch_dtype) + + def _init_pipeline(self, *args, **kwargs): + transformer = FluxTransformer2DModel.from_single_file( + self.gguf_ckpt, quantization_config=self.quantization_config, torch_dtype=self.torch_dtype + ) + pipe = DiffusionPipeline.from_pretrained( + "black-forest-labs/FLUX.1-dev", transformer=transformer, torch_dtype=self.torch_dtype + ) + return pipe + + +# ======================== NVIDIA ModelOpt ======================== + + +# Model-level ModelOpt tests live in `tests/models/testing_utils/quantization.py` +# (`ModelOptTesterMixin` / `ModelOptCompileTesterMixin`), wired into model test files via concrete +# classes (e.g. `TestSD3TransformerModelOpt`). Only pipeline-level coverage remains here. +@is_quantization +@is_modelopt +@nightly +@require_big_accelerator +@require_accelerate +@require_modelopt_version_greater_or_equal("0.33.1") +class TestModelOptFP8: + model_id = "hf-internal-testing/tiny-sd3-pipe" + + @pytest.fixture(autouse=True) + def _setup(self): + backend_reset_peak_memory_stats(torch_device) + backend_empty_cache(torch_device) + gc.collect() + yield + backend_reset_peak_memory_stats(torch_device) + backend_empty_cache(torch_device) + gc.collect() + + def test_model_cpu_offload(self): + transformer = SD3Transformer2DModel.from_pretrained( + self.model_id, + quantization_config=NVIDIAModelOptConfig(quant_type="FP8"), + subfolder="transformer", + torch_dtype=torch.bfloat16, + ) + pipe = StableDiffusion3Pipeline.from_pretrained( + self.model_id, transformer=transformer, torch_dtype=torch.bfloat16 + ) + pipe.enable_model_cpu_offload(device=torch_device) + _ = pipe("a cat holding a sign that says hello", num_inference_steps=2) diff --git a/tests/quantization/bnb/README.md b/tests/quantization/bnb/README.md deleted file mode 100644 index c4bdd53b765e..000000000000 --- a/tests/quantization/bnb/README.md +++ /dev/null @@ -1,41 +0,0 @@ -The tests here are adapted from [`transformers` tests](https://github.com/huggingface/transformers/tree/409fcfdfccde77a14b7cc36972b774cabc371ae1/tests/quantization/bnb). - -They were conducted on the `audace` machine, using a single RTX 4090. Below is `nvidia-smi`: - -```bash -+-----------------------------------------------------------------------------------------+ -| NVIDIA-SMI 550.90.07 Driver Version: 550.90.07 CUDA Version: 12.4 | -|-----------------------------------------+------------------------+----------------------+ -| GPU Name Persistence-M | Bus-Id Disp.A | Volatile Uncorr. ECC | -| Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. | -| | | MIG M. | -|=========================================+========================+======================| -| 0 NVIDIA GeForce RTX 4090 Off | 00000000:01:00.0 Off | Off | -| 30% 55C P0 61W / 450W | 1MiB / 24564MiB | 2% Default | -| | | N/A | -+-----------------------------------------+------------------------+----------------------+ -| 1 NVIDIA GeForce RTX 4090 Off | 00000000:13:00.0 Off | Off | -| 30% 51C P0 60W / 450W | 1MiB / 24564MiB | 0% Default | -| | | N/A | -+-----------------------------------------+------------------------+----------------------+ -``` - -`diffusers-cli`: - -```bash -- 🤗 Diffusers version: 0.31.0.dev0 -- Platform: Linux-5.15.0-117-generic-x86_64-with-glibc2.35 -- Running on Google Colab?: No -- Python version: 3.10.12 -- PyTorch version (GPU?): 2.5.0.dev20240818+cu124 (True) -- Huggingface_hub version: 0.24.5 -- Transformers version: 4.44.2 -- Accelerate version: 0.34.0.dev0 -- PEFT version: 0.12.0 -- Bitsandbytes version: 0.43.3 -- Safetensors version: 0.4.4 -- xFormers version: not installed -- Accelerator: NVIDIA GeForce RTX 4090, 24564 MiB -NVIDIA GeForce RTX 4090, 24564 MiB -- Using GPU in script?: Yes -``` \ No newline at end of file diff --git a/tests/quantization/bnb/test_4bit.py b/tests/quantization/bnb/test_4bit.py index cf3671439efb..f666c5899f9b 100644 --- a/tests/quantization/bnb/test_4bit.py +++ b/tests/quantization/bnb/test_4bit.py @@ -1,9 +1,9 @@ # coding=utf-8 -# Copyright 2025 The HuggingFace Team Inc. +# Copyright 2026 The HuggingFace Team 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 clone of the License at +# You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # @@ -13,192 +13,53 @@ # See the License for the specific language governing permissions and # limitations under the License. import gc -import os -import tempfile -import numpy as np import pytest -import safetensors.torch -from huggingface_hub import hf_hub_download -from PIL import Image +import torch -from diffusers import ( - BitsAndBytesConfig, - DiffusionPipeline, - FluxControlPipeline, - FluxTransformer2DModel, - SD3Transformer2DModel, -) -from diffusers.quantizers import PipelineQuantizationConfig -from diffusers.utils import is_accelerate_version, logging +from diffusers import BitsAndBytesConfig +from diffusers.utils import logging from ...testing_utils import ( CaptureLogger, backend_empty_cache, + is_bitsandbytes, is_bitsandbytes_available, - is_torch_available, - is_transformers_available, - load_pt, - numpy_cosine_similarity_distance, + is_quantization, require_accelerate, require_bitsandbytes_version_greater, - require_peft_backend, require_torch, require_torch_accelerator, - require_torch_version_greater, - require_transformers_version_greater, slow, torch_device, ) -from ..test_torch_compile_utils import QuantCompileTests - - -def get_some_linear_layer(model): - if model.__class__.__name__ in ["SD3Transformer2DModel", "FluxTransformer2DModel"]: - return model.transformer_blocks[0].attn.to_q - else: - return NotImplementedError("Don't know what layer to retrieve here.") - - -if is_transformers_available(): - from transformers import BitsAndBytesConfig as BnbConfig - from transformers import T5EncoderModel - -if is_torch_available(): - import torch - - from ..utils import get_memory_consumption_stat if is_bitsandbytes_available(): - import bitsandbytes as bnb - from diffusers.quantizers.bitsandbytes.utils import replace_with_bnb_linear +# Model-level BitsAndBytes tests live in `tests/models/testing_utils/quantization.py` and +# pipeline-level ones in `tests/pipelines/testing_utils/quantization.py`. This module covers +# backend behavior that fits neither: config validation, loading error paths, and utility warnings. +@is_quantization +@is_bitsandbytes @require_bitsandbytes_version_greater("0.43.2") @require_accelerate @require_torch @require_torch_accelerator @slow -class Base4bitTests: - # We need to test on relatively large models (aka >1b parameters otherwise the quantiztion may not work as expected) - # Therefore here we use only SD3 to test our module +class TestBnB4BitBasic: model_name = "stabilityai/stable-diffusion-3-medium-diffusers" - # This was obtained on audace so the number might slightly change - expected_rel_difference = 3.69 - - expected_memory_saving_ratio = 0.8 - - prompt = "a beautiful sunset amidst the mountains." - num_inference_steps = 10 - seed = 0 - - @pytest.fixture(autouse=True, scope="class") - def _toggle_determinism(self): - was_enabled = torch.are_deterministic_algorithms_enabled() - if not was_enabled: - torch.use_deterministic_algorithms(True) - yield - if not was_enabled: - torch.use_deterministic_algorithms(False) - - def get_dummy_inputs(self): - prompt_embeds = load_pt( - "https://huggingface.co/datasets/hf-internal-testing/bnb-diffusers-testing-artifacts/resolve/main/prompt_embeds.pt", - torch_device, - ) - pooled_prompt_embeds = load_pt( - "https://huggingface.co/datasets/hf-internal-testing/bnb-diffusers-testing-artifacts/resolve/main/pooled_prompt_embeds.pt", - torch_device, - ) - latent_model_input = load_pt( - "https://huggingface.co/datasets/hf-internal-testing/bnb-diffusers-testing-artifacts/resolve/main/latent_model_input.pt", - torch_device, - ) - - input_dict_for_transformer = { - "hidden_states": latent_model_input, - "encoder_hidden_states": prompt_embeds, - "pooled_projections": pooled_prompt_embeds, - "timestep": torch.Tensor([1.0]), - "return_dict": False, - } - return input_dict_for_transformer - - -class TestBnB4BitBasic(Base4bitTests): @pytest.fixture(autouse=True) def _setup_basic(self): gc.collect() backend_empty_cache(torch_device) - - # Models - self.model_fp16 = SD3Transformer2DModel.from_pretrained( - self.model_name, subfolder="transformer", torch_dtype=torch.float16 - ) - nf4_config = BitsAndBytesConfig( - load_in_4bit=True, - bnb_4bit_quant_type="nf4", - bnb_4bit_compute_dtype=torch.float16, - ) - self.model_4bit = SD3Transformer2DModel.from_pretrained( - self.model_name, subfolder="transformer", quantization_config=nf4_config, device_map=torch_device - ) yield - if hasattr(self, "model_fp16"): - del self.model_fp16 - if hasattr(self, "model_4bit"): - del self.model_4bit - gc.collect() backend_empty_cache(torch_device) - def test_model_memory_usage(self): - # Delete to not let anything interfere. - del self.model_4bit, self.model_fp16 - - # Re-instantiate. - inputs = self.get_dummy_inputs() - inputs = { - k: v.to(device=torch_device, dtype=torch.float16) for k, v in inputs.items() if not isinstance(v, bool) - } - model_fp16 = SD3Transformer2DModel.from_pretrained( - self.model_name, subfolder="transformer", torch_dtype=torch.float16 - ).to(torch_device) - unquantized_model_memory = get_memory_consumption_stat(model_fp16, inputs) - del model_fp16 - - nf4_config = BitsAndBytesConfig( - load_in_4bit=True, - bnb_4bit_quant_type="nf4", - bnb_4bit_compute_dtype=torch.float16, - ) - model_4bit = SD3Transformer2DModel.from_pretrained( - self.model_name, subfolder="transformer", quantization_config=nf4_config, torch_dtype=torch.float16 - ) - quantized_model_memory = get_memory_consumption_stat(model_4bit, inputs) - assert unquantized_model_memory / quantized_model_memory >= self.expected_memory_saving_ratio - - def test_device_assignment(self): - mem_before = self.model_4bit.get_memory_footprint() - - # Move to CPU - self.model_4bit.to("cpu") - assert self.model_4bit.device.type == "cpu" - assert self.model_4bit.get_memory_footprint() == pytest.approx(mem_before) - - # Move back to CUDA device - for device in [0, f"{torch_device}", f"{torch_device}:0", "call()"]: - if device == "call()": - self.model_4bit.to(f"{torch_device}:0") - else: - self.model_4bit.to(device) - assert self.model_4bit.device == torch.device(0) - assert self.model_4bit.get_memory_footprint() == pytest.approx(mem_before) - self.model_4bit.to("cpu") - def test_bnb_4bit_wrong_config(self): r""" Test whether creating a bnb config with unsupported values leads to errors. @@ -206,36 +67,6 @@ def test_bnb_4bit_wrong_config(self): with pytest.raises(ValueError): _ = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_storage="add") - def test_bnb_4bit_errors_loading_incorrect_state_dict(self): - r""" - Test if loading with an incorrect state dict raises an error. - """ - with tempfile.TemporaryDirectory() as tmpdirname: - nf4_config = BitsAndBytesConfig(load_in_4bit=True) - model_4bit = SD3Transformer2DModel.from_pretrained( - self.model_name, subfolder="transformer", quantization_config=nf4_config, device_map=torch_device - ) - model_4bit.save_pretrained(tmpdirname) - del model_4bit - - with pytest.raises(ValueError) as err_context: - state_dict = safetensors.torch.load_file( - os.path.join(tmpdirname, "diffusion_pytorch_model.safetensors") - ) - - # corrupt the state dict - key_to_target = "context_embedder.weight" # can be other keys too. - compatible_param = state_dict[key_to_target] - corrupted_param = torch.randn(compatible_param.shape[0] - 1, 1) - state_dict[key_to_target] = bnb.nn.Params4bit(corrupted_param, requires_grad=False) - safetensors.torch.save_file( - state_dict, os.path.join(tmpdirname, "diffusion_pytorch_model.safetensors") - ) - - _ = SD3Transformer2DModel.from_pretrained(tmpdirname) - - assert key_to_target in str(err_context.value) - def test_bnb_4bit_logs_warning_for_no_quantization(self): model_with_no_linear = torch.nn.Sequential(torch.nn.Conv2d(4, 4, 3), torch.nn.ReLU()) quantization_config = BitsAndBytesConfig(load_in_4bit=True) @@ -247,441 +78,3 @@ def test_bnb_4bit_logs_warning_for_no_quantization(self): "You are loading your model in 8bit or 4bit but no linear modules were found in your model." in cap_logger.out ) - - -@require_transformers_version_greater("4.44.0") -class TestSlowBnb4Bit(Base4bitTests): - @pytest.fixture(autouse=True) - def _setup_slow(self): - gc.collect() - backend_empty_cache(torch_device) - - nf4_config = BitsAndBytesConfig( - load_in_4bit=True, - bnb_4bit_quant_type="nf4", - bnb_4bit_compute_dtype=torch.float16, - ) - model_4bit = SD3Transformer2DModel.from_pretrained( - self.model_name, subfolder="transformer", quantization_config=nf4_config, device_map=torch_device - ) - self.pipeline_4bit = DiffusionPipeline.from_pretrained( - self.model_name, transformer=model_4bit, torch_dtype=torch.float16 - ) - self.pipeline_4bit.enable_model_cpu_offload() - yield - del self.pipeline_4bit - - gc.collect() - backend_empty_cache(torch_device) - - def test_quality(self): - output = self.pipeline_4bit( - prompt=self.prompt, - num_inference_steps=self.num_inference_steps, - generator=torch.manual_seed(self.seed), - output_type="np", - ).images - - out_slice = output[0, -3:, -3:, -1].flatten() - expected_slice = np.array([0.1123, 0.1296, 0.1609, 0.1042, 0.1230, 0.1274, 0.0928, 0.1165, 0.1216]) - - max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice) - assert max_diff < 1e-2 - - def test_generate_quality_dequantize(self): - r""" - Test that loading the model and unquantize it produce correct results. - """ - self.pipeline_4bit.transformer.dequantize() - output = self.pipeline_4bit( - prompt=self.prompt, - num_inference_steps=self.num_inference_steps, - generator=torch.manual_seed(self.seed), - output_type="np", - ).images - - out_slice = output[0, -3:, -3:, -1].flatten() - expected_slice = np.array([0.1216, 0.1387, 0.1584, 0.1152, 0.1318, 0.1282, 0.1062, 0.1226, 0.1228]) - max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice) - assert max_diff < 1e-3 - - # Since we offloaded the `pipeline_4bit.transformer` to CPU (result of `enable_model_cpu_offload()), check - # the following. - assert self.pipeline_4bit.transformer.device.type == "cpu" - # calling it again shouldn't be a problem - _ = self.pipeline_4bit( - prompt=self.prompt, - num_inference_steps=2, - generator=torch.manual_seed(self.seed), - output_type="np", - ).images - - def test_moving_to_cpu_throws_warning(self): - nf4_config = BitsAndBytesConfig( - load_in_4bit=True, - bnb_4bit_quant_type="nf4", - bnb_4bit_compute_dtype=torch.float16, - ) - model_4bit = SD3Transformer2DModel.from_pretrained( - self.model_name, subfolder="transformer", quantization_config=nf4_config, device_map=torch_device - ) - - logger = logging.get_logger("diffusers.pipelines.pipeline_utils") - logger.setLevel(30) - with CaptureLogger(logger) as cap_logger: - # Because `model.dtype` will return torch.float16 as SD3 transformer has - # a conv layer as the first layer. - _ = DiffusionPipeline.from_pretrained( - self.model_name, transformer=model_4bit, torch_dtype=torch.float16 - ).to("cpu") - - assert "Pipelines loaded with `dtype=torch.float16`" in cap_logger.out - - @pytest.mark.xfail( - condition=is_accelerate_version("<=", "1.1.1"), - reason="Test will pass after https://github.com/huggingface/accelerate/pull/3223 is in a release.", - strict=True, - ) - def test_pipeline_cuda_placement_works_with_nf4(self): - transformer_nf4_config = BitsAndBytesConfig( - load_in_4bit=True, - bnb_4bit_quant_type="nf4", - bnb_4bit_compute_dtype=torch.float16, - ) - transformer_4bit = SD3Transformer2DModel.from_pretrained( - self.model_name, - subfolder="transformer", - quantization_config=transformer_nf4_config, - torch_dtype=torch.float16, - device_map=torch_device, - ) - text_encoder_3_nf4_config = BnbConfig( - load_in_4bit=True, - bnb_4bit_quant_type="nf4", - bnb_4bit_compute_dtype=torch.float16, - ) - text_encoder_3_4bit = T5EncoderModel.from_pretrained( - self.model_name, - subfolder="text_encoder_3", - quantization_config=text_encoder_3_nf4_config, - torch_dtype=torch.float16, - device_map=torch_device, - ) - # CUDA device placement works. - pipeline_4bit = DiffusionPipeline.from_pretrained( - self.model_name, - transformer=transformer_4bit, - text_encoder_3=text_encoder_3_4bit, - torch_dtype=torch.float16, - ).to(torch_device) - - # Check if inference works. - _ = pipeline_4bit(self.prompt, max_sequence_length=20, num_inference_steps=2) - - del pipeline_4bit - - def test_device_map(self): - """ - Test if the quantized model is working properly with "auto". - cpu/disk offloading as well doesn't work with bnb. - """ - - def get_dummy_tensor_inputs(device=None, seed: int = 0): - batch_size = 1 - num_latent_channels = 4 - num_image_channels = 3 - height = width = 4 - sequence_length = 48 - embedding_dim = 32 - - torch.manual_seed(seed) - hidden_states = torch.randn((batch_size, height * width, num_latent_channels)).to( - device, dtype=torch.bfloat16 - ) - torch.manual_seed(seed) - encoder_hidden_states = torch.randn((batch_size, sequence_length, embedding_dim)).to( - device, dtype=torch.bfloat16 - ) - - torch.manual_seed(seed) - pooled_prompt_embeds = torch.randn((batch_size, embedding_dim)).to(device, dtype=torch.bfloat16) - - torch.manual_seed(seed) - text_ids = torch.randn((sequence_length, num_image_channels)).to(device, dtype=torch.bfloat16) - - torch.manual_seed(seed) - image_ids = torch.randn((height * width, num_image_channels)).to(device, dtype=torch.bfloat16) - - timestep = torch.tensor([1.0]).to(device, dtype=torch.bfloat16).expand(batch_size) - - return { - "hidden_states": hidden_states, - "encoder_hidden_states": encoder_hidden_states, - "pooled_projections": pooled_prompt_embeds, - "txt_ids": text_ids, - "img_ids": image_ids, - "timestep": timestep, - } - - inputs = get_dummy_tensor_inputs(torch_device) - expected_slice = np.array( - [0.47070312, 0.00390625, -0.03662109, -0.19628906, -0.53125, 0.5234375, -0.17089844, -0.59375, 0.578125] - ) - - # non sharded - quantization_config = BitsAndBytesConfig( - load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.float16 - ) - quantized_model = FluxTransformer2DModel.from_pretrained( - "hf-internal-testing/tiny-flux-pipe", - subfolder="transformer", - quantization_config=quantization_config, - device_map="auto", - torch_dtype=torch.bfloat16, - ) - - weight = quantized_model.transformer_blocks[0].ff.net[2].weight - assert isinstance(weight, bnb.nn.modules.Params4bit) - - output = quantized_model(**inputs)[0] - output_slice = output.flatten()[-9:].detach().float().cpu().numpy() - assert numpy_cosine_similarity_distance(output_slice, expected_slice) < 1e-3 - - # sharded - - quantization_config = BitsAndBytesConfig( - load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.float16 - ) - quantized_model = FluxTransformer2DModel.from_pretrained( - "hf-internal-testing/tiny-flux-sharded", - subfolder="transformer", - quantization_config=quantization_config, - device_map="auto", - torch_dtype=torch.bfloat16, - ) - - weight = quantized_model.transformer_blocks[0].ff.net[2].weight - assert isinstance(weight, bnb.nn.modules.Params4bit) - - output = quantized_model(**inputs)[0] - output_slice = output.flatten()[-9:].detach().float().cpu().numpy() - - assert numpy_cosine_similarity_distance(output_slice, expected_slice) < 1e-3 - - -@require_transformers_version_greater("4.44.0") -class TestSlowBnb4BitFlux(Base4bitTests): - @pytest.fixture(autouse=True) - def _setup_flux(self): - gc.collect() - backend_empty_cache(torch_device) - - model_id = "hf-internal-testing/flux.1-dev-nf4-pkg" - t5_4bit = T5EncoderModel.from_pretrained(model_id, subfolder="text_encoder_2") - transformer_4bit = FluxTransformer2DModel.from_pretrained(model_id, subfolder="transformer") - self.pipeline_4bit = DiffusionPipeline.from_pretrained( - "black-forest-labs/FLUX.1-dev", - text_encoder_2=t5_4bit, - transformer=transformer_4bit, - torch_dtype=torch.float16, - ) - self.pipeline_4bit.enable_model_cpu_offload() - yield - del self.pipeline_4bit - - gc.collect() - backend_empty_cache(torch_device) - - def test_quality(self): - # keep the resolution and max tokens to a lower number for faster execution. - output = self.pipeline_4bit( - prompt=self.prompt, - num_inference_steps=self.num_inference_steps, - generator=torch.manual_seed(self.seed), - height=256, - width=256, - max_sequence_length=64, - output_type="np", - ).images - - out_slice = output[0, -3:, -3:, -1].flatten() - expected_slice = np.array([0.0583, 0.0586, 0.0632, 0.0815, 0.0813, 0.0947, 0.1040, 0.1145, 0.1265]) - - max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice) - assert max_diff < 1e-3 - - @require_peft_backend - def test_lora_loading(self): - self.pipeline_4bit.load_lora_weights( - hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors"), adapter_name="hyper-sd" - ) - self.pipeline_4bit.set_adapters("hyper-sd", adapter_weights=0.125) - - output = self.pipeline_4bit( - prompt=self.prompt, - height=256, - width=256, - max_sequence_length=64, - output_type="np", - num_inference_steps=8, - generator=torch.Generator().manual_seed(42), - ).images - out_slice = output[0, -3:, -3:, -1].flatten() - expected_slice = np.array([0.5347, 0.5342, 0.5283, 0.5093, 0.4988, 0.5093, 0.5044, 0.5015, 0.4946]) - - max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice) - assert max_diff < 1e-3 - - -@require_transformers_version_greater("4.44.0") -@require_peft_backend -class TestSlowBnb4BitFluxControlWithLora(Base4bitTests): - @pytest.fixture(autouse=True) - def _setup_flux_control(self): - gc.collect() - backend_empty_cache(torch_device) - - self.pipeline_4bit = FluxControlPipeline.from_pretrained("eramth/flux-4bit", torch_dtype=torch.float16) - self.pipeline_4bit.enable_model_cpu_offload() - yield - del self.pipeline_4bit - - gc.collect() - backend_empty_cache(torch_device) - - def test_lora_loading(self): - self.pipeline_4bit.load_lora_weights("black-forest-labs/FLUX.1-Canny-dev-lora") - - output = self.pipeline_4bit( - prompt=self.prompt, - control_image=Image.new(mode="RGB", size=(256, 256)), - height=256, - width=256, - max_sequence_length=64, - output_type="np", - num_inference_steps=8, - generator=torch.Generator().manual_seed(42), - ).images - out_slice = output[0, -3:, -3:, -1].flatten() - expected_slice = np.array([0.1636, 0.1675, 0.1982, 0.1743, 0.1809, 0.1936, 0.1743, 0.2095, 0.2139]) - - max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice) - assert max_diff < 1e-3, f"{out_slice=} != {expected_slice=}" - - -@slow -class TestBnb4BitSerialization(Base4bitTests): - @pytest.fixture(autouse=True) - def _setup_serialization(self): - yield - gc.collect() - backend_empty_cache(torch_device) - - def test_serialization(self, quant_type="nf4", double_quant=True, safe_serialization=True): - r""" - Test whether it is possible to serialize a model in 4-bit. Uses most typical params as default. - See ExtendedSerializationTest class for more params combinations. - """ - - self.quantization_config = BitsAndBytesConfig( - load_in_4bit=True, - bnb_4bit_quant_type=quant_type, - bnb_4bit_use_double_quant=double_quant, - bnb_4bit_compute_dtype=torch.bfloat16, - ) - model_0 = SD3Transformer2DModel.from_pretrained( - self.model_name, - subfolder="transformer", - quantization_config=self.quantization_config, - device_map=torch_device, - ) - assert "_pre_quantization_dtype" in model_0.config - with tempfile.TemporaryDirectory() as tmpdirname: - model_0.save_pretrained(tmpdirname, safe_serialization=safe_serialization) - - config = SD3Transformer2DModel.load_config(tmpdirname) - assert "quantization_config" in config - assert "_pre_quantization_dtype" not in config - - model_1 = SD3Transformer2DModel.from_pretrained(tmpdirname) - - # checking quantized linear module weight - linear = get_some_linear_layer(model_1) - assert linear.weight.__class__ == bnb.nn.Params4bit - assert hasattr(linear.weight, "quant_state") - assert linear.weight.quant_state.__class__ == bnb.functional.QuantState - - # checking memory footpring - assert model_0.get_memory_footprint() / model_1.get_memory_footprint() == pytest.approx(1, abs=10**-2) - - # Matching all parameters and their quant_state items: - d0 = dict(model_0.named_parameters()) - d1 = dict(model_1.named_parameters()) - assert d0.keys() == d1.keys() - - for k in d0.keys(): - assert d0[k].shape == d1[k].shape - assert d0[k].device.type == d1[k].device.type - assert d0[k].device == d1[k].device - assert d0[k].dtype == d1[k].dtype - assert torch.equal(d0[k], d1[k].to(d0[k].device)) - - if isinstance(d0[k], bnb.nn.modules.Params4bit): - for v0, v1 in zip( - d0[k].quant_state.as_dict().values(), - d1[k].quant_state.as_dict().values(), - ): - if isinstance(v0, torch.Tensor): - assert torch.equal(v0, v1.to(v0.device)) - else: - assert v0 == v1 - - # comparing forward() outputs - dummy_inputs = self.get_dummy_inputs() - inputs = {k: v.to(torch_device) for k, v in dummy_inputs.items() if isinstance(v, torch.Tensor)} - inputs.update({k: v for k, v in dummy_inputs.items() if k not in inputs}) - out_0 = model_0(**inputs)[0] - out_1 = model_1(**inputs)[0] - assert torch.equal(out_0, out_1) - - -class TestExtendedSerialization(TestBnb4BitSerialization): - """ - tests more combinations of parameters - """ - - def test_nf4_single_unsafe(self): - self.test_serialization(quant_type="nf4", double_quant=False, safe_serialization=False) - - def test_nf4_double_unsafe(self): - self.test_serialization(quant_type="nf4", double_quant=True, safe_serialization=False) - - # nf4 double safetensors quantization is tested in test_serialization() method from the parent class - - def test_fp4_single_unsafe(self): - self.test_serialization(quant_type="fp4", double_quant=False, safe_serialization=False) - - def test_fp4_single_safe(self): - self.test_serialization(quant_type="fp4", double_quant=False, safe_serialization=True) - - def test_fp4_double_unsafe(self): - self.test_serialization(quant_type="fp4", double_quant=True, safe_serialization=False) - - def test_fp4_double_safe(self): - self.test_serialization(quant_type="fp4", double_quant=True, safe_serialization=True) - - -@require_torch_version_greater("2.7.1") -@require_bitsandbytes_version_greater("0.45.5") -class TestBnb4BitCompile(QuantCompileTests): - @property - def quantization_config(self): - return PipelineQuantizationConfig( - quant_backend="bitsandbytes_4bit", - quant_kwargs={ - "load_in_4bit": True, - "bnb_4bit_quant_type": "nf4", - "bnb_4bit_compute_dtype": torch.bfloat16, - }, - components_to_quantize=["transformer", "text_encoder_2"], - ) diff --git a/tests/quantization/bnb/test_mixed_int8.py b/tests/quantization/bnb/test_mixed_int8.py index 5100e5fc353b..d37099ff623c 100644 --- a/tests/quantization/bnb/test_mixed_int8.py +++ b/tests/quantization/bnb/test_mixed_int8.py @@ -1,9 +1,9 @@ # coding=utf-8 -# Copyright 2025 The HuggingFace Team Inc. +# Copyright 2026 The HuggingFace Team 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 clone of the License at +# You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # @@ -13,271 +13,51 @@ # See the License for the specific language governing permissions and # limitations under the License. import gc -import tempfile -import numpy as np import pytest -from huggingface_hub import hf_hub_download -from PIL import Image +import torch -from diffusers import ( - BitsAndBytesConfig, - DiffusionPipeline, - FluxControlPipeline, - FluxTransformer2DModel, - SanaTransformer2DModel, - SD3Transformer2DModel, - logging, -) -from diffusers.quantizers import PipelineQuantizationConfig -from diffusers.utils import is_accelerate_version +from diffusers import BitsAndBytesConfig +from diffusers.utils import logging from ...testing_utils import ( CaptureLogger, - Expectations, backend_empty_cache, + is_bitsandbytes, is_bitsandbytes_available, - is_torch_available, - is_transformers_available, - load_pt, - numpy_cosine_similarity_distance, + is_quantization, require_accelerate, - require_big_accelerator, require_bitsandbytes_version_greater, - require_peft_backend, - require_peft_version_greater, require_torch, require_torch_accelerator, - require_torch_version_greater_equal, - require_transformers_version_greater, slow, torch_device, ) -from ..test_torch_compile_utils import QuantCompileTests - - -def get_some_linear_layer(model): - if model.__class__.__name__ in ["SD3Transformer2DModel", "FluxTransformer2DModel"]: - return model.transformer_blocks[0].attn.to_q - else: - return NotImplementedError("Don't know what layer to retrieve here.") - - -if is_transformers_available(): - from transformers import BitsAndBytesConfig as BnbConfig - from transformers import T5EncoderModel - -if is_torch_available(): - import torch - - from ..utils import LoRALayer, get_memory_consumption_stat if is_bitsandbytes_available(): - import bitsandbytes as bnb - - from diffusers.quantizers.bitsandbytes import replace_with_bnb_linear + from diffusers.quantizers.bitsandbytes.utils import replace_with_bnb_linear +# Model-level BitsAndBytes tests live in `tests/models/testing_utils/quantization.py` and +# pipeline-level ones in `tests/pipelines/testing_utils/quantization.py`. This module covers +# backend behavior that fits neither: utility warnings. +@is_quantization +@is_bitsandbytes @require_bitsandbytes_version_greater("0.43.2") @require_accelerate @require_torch @require_torch_accelerator @slow -class Base8bitTests: - # We need to test on relatively large models (aka >1b parameters otherwise the quantiztion may not work as expected) - # Therefore here we use only SD3 to test our module - model_name = "stabilityai/stable-diffusion-3-medium-diffusers" - - # This was obtained on audace so the number might slightly change - expected_rel_difference = 1.94 - - expected_memory_saving_ratio = 0.7 - - prompt = "a beautiful sunset amidst the mountains." - num_inference_steps = 10 - seed = 0 - - @pytest.fixture(autouse=True, scope="class") - def _toggle_determinism(self): - was_enabled = torch.are_deterministic_algorithms_enabled() - if not was_enabled: - torch.use_deterministic_algorithms(True) - yield - if not was_enabled: - torch.use_deterministic_algorithms(False) - - def get_dummy_inputs(self): - prompt_embeds = load_pt( - "https://huggingface.co/datasets/hf-internal-testing/bnb-diffusers-testing-artifacts/resolve/main/prompt_embeds.pt", - map_location="cpu", - ) - pooled_prompt_embeds = load_pt( - "https://huggingface.co/datasets/hf-internal-testing/bnb-diffusers-testing-artifacts/resolve/main/pooled_prompt_embeds.pt", - map_location="cpu", - ) - latent_model_input = load_pt( - "https://huggingface.co/datasets/hf-internal-testing/bnb-diffusers-testing-artifacts/resolve/main/latent_model_input.pt", - map_location="cpu", - ) - - input_dict_for_transformer = { - "hidden_states": latent_model_input, - "encoder_hidden_states": prompt_embeds, - "pooled_projections": pooled_prompt_embeds, - "timestep": torch.Tensor([1.0]), - "return_dict": False, - } - return input_dict_for_transformer - - -class TestBnB8bitBasic(Base8bitTests): +class TestBnB8bitBasic: @pytest.fixture(autouse=True) def _setup_basic(self): gc.collect() backend_empty_cache(torch_device) - - # Models - self.model_fp16 = SD3Transformer2DModel.from_pretrained( - self.model_name, subfolder="transformer", torch_dtype=torch.float16 - ) - mixed_int8_config = BitsAndBytesConfig(load_in_8bit=True) - self.model_8bit = SD3Transformer2DModel.from_pretrained( - self.model_name, subfolder="transformer", quantization_config=mixed_int8_config, device_map=torch_device - ) yield - if hasattr(self, "model_fp16"): - del self.model_fp16 - if hasattr(self, "model_8bit"): - del self.model_8bit - gc.collect() backend_empty_cache(torch_device) - def test_model_memory_usage(self): - # Delete to not let anything interfere. - del self.model_8bit, self.model_fp16 - - # Re-instantiate. - inputs = self.get_dummy_inputs() - inputs = { - k: v.to(device=torch_device, dtype=torch.float16) for k, v in inputs.items() if not isinstance(v, bool) - } - model_fp16 = SD3Transformer2DModel.from_pretrained( - self.model_name, subfolder="transformer", torch_dtype=torch.float16 - ).to(torch_device) - unquantized_model_memory = get_memory_consumption_stat(model_fp16, inputs) - del model_fp16 - - config = BitsAndBytesConfig(load_in_8bit=True) - model_8bit = SD3Transformer2DModel.from_pretrained( - self.model_name, subfolder="transformer", quantization_config=config, torch_dtype=torch.float16 - ) - quantized_model_memory = get_memory_consumption_stat(model_8bit, inputs) - assert unquantized_model_memory / quantized_model_memory >= self.expected_memory_saving_ratio - - def test_original_dtype(self): - r""" - A simple test to check if the model successfully stores the original dtype - """ - assert "_pre_quantization_dtype" in self.model_8bit.config - assert "_pre_quantization_dtype" not in self.model_fp16.config - assert self.model_8bit.config["_pre_quantization_dtype"] == torch.float16 - - def test_keep_modules_in_fp32(self): - r""" - A simple tests to check if the modules under `_keep_in_fp32_modules` are kept in fp32. - Also ensures if inference works. - """ - fp32_modules = SD3Transformer2DModel._keep_in_fp32_modules - SD3Transformer2DModel._keep_in_fp32_modules = ["proj_out"] - - mixed_int8_config = BitsAndBytesConfig(load_in_8bit=True) - model = SD3Transformer2DModel.from_pretrained( - self.model_name, subfolder="transformer", quantization_config=mixed_int8_config, device_map=torch_device - ) - - for name, module in model.named_modules(): - if isinstance(module, torch.nn.Linear): - if name in model._keep_in_fp32_modules: - assert module.weight.dtype == torch.float32 - else: - # 8-bit parameters are packed in int8 variables - assert module.weight.dtype == torch.int8 - - # test if inference works. - with torch.no_grad() and torch.autocast(model.device.type, dtype=torch.float16): - input_dict_for_transformer = self.get_dummy_inputs() - model_inputs = { - k: v.to(device=torch_device) for k, v in input_dict_for_transformer.items() if not isinstance(v, bool) - } - model_inputs.update({k: v for k, v in input_dict_for_transformer.items() if k not in model_inputs}) - _ = model(**model_inputs) - - SD3Transformer2DModel._keep_in_fp32_modules = fp32_modules - - def test_llm_skip(self): - r""" - A simple test to check if `llm_int8_skip_modules` works as expected - """ - config = BitsAndBytesConfig(load_in_8bit=True, llm_int8_skip_modules=["proj_out"]) - model_8bit = SD3Transformer2DModel.from_pretrained( - self.model_name, subfolder="transformer", quantization_config=config, device_map=torch_device - ) - linear = get_some_linear_layer(model_8bit) - assert linear.weight.dtype == torch.int8 - assert isinstance(linear, bnb.nn.Linear8bitLt) - - assert isinstance(model_8bit.proj_out, torch.nn.Linear) - assert model_8bit.proj_out.weight.dtype != torch.int8 - - @require_bitsandbytes_version_greater("0.48.0") - def test_device_and_dtype_assignment(self): - r""" - Test whether trying to cast (or assigning a device to) a model after converting it in 8-bit will throw an error. - Checks also if other models are casted correctly. - """ - - with pytest.raises(ValueError): - # Tries with a `dtype`` - self.model_8bit.to(torch.float16) - - with pytest.raises(ValueError): - # Tries with a `device` - self.model_8bit.float() - - with pytest.raises(ValueError): - # Tries with a `dtype` - self.model_8bit.half() - - # This should work with 0.48.0 - self.model_8bit.to("cpu") - self.model_8bit.to(torch.device(f"{torch_device}:0")) - - # Test if we did not break anything - self.model_fp16 = self.model_fp16.to(dtype=torch.float32, device=torch_device) - input_dict_for_transformer = self.get_dummy_inputs() - model_inputs = { - k: v.to(dtype=torch.float32, device=torch_device) - for k, v in input_dict_for_transformer.items() - if not isinstance(v, bool) - } - model_inputs.update({k: v for k, v in input_dict_for_transformer.items() if k not in model_inputs}) - with torch.no_grad(): - _ = self.model_fp16(**model_inputs) - - # Check this does not throw an error - _ = self.model_fp16.to("cpu") - - # Check this does not throw an error - _ = self.model_fp16.half() - - # Check this does not throw an error - _ = self.model_fp16.float() - - # Check that this does not throw an error - _ = self.model_fp16.to(torch_device) - def test_bnb_8bit_logs_warning_for_no_quantization(self): model_with_no_linear = torch.nn.Sequential(torch.nn.Conv2d(4, 4, 3), torch.nn.ReLU()) quantization_config = BitsAndBytesConfig(load_in_8bit=True) @@ -289,514 +69,3 @@ def test_bnb_8bit_logs_warning_for_no_quantization(self): "You are loading your model in 8bit or 4bit but no linear modules were found in your model." in cap_logger.out ) - - -class TestBnb8bitDevice(Base8bitTests): - @pytest.fixture(autouse=True) - def _setup_device(self): - gc.collect() - backend_empty_cache(torch_device) - - mixed_int8_config = BitsAndBytesConfig(load_in_8bit=True) - self.model_8bit = SanaTransformer2DModel.from_pretrained( - "Efficient-Large-Model/Sana_1600M_4Kpx_BF16_diffusers", - subfolder="transformer", - quantization_config=mixed_int8_config, - device_map=torch_device, - ) - yield - del self.model_8bit - - gc.collect() - backend_empty_cache(torch_device) - - def test_buffers_device_assignment(self): - for buffer_name, buffer in self.model_8bit.named_buffers(): - assert buffer.device.type == torch.device(torch_device).type, ( - f"Expected device {torch_device} for {buffer_name} got {buffer.device}." - ) - - -class TestBnB8bitTraining(Base8bitTests): - @pytest.fixture(autouse=True) - def _setup_training(self): - gc.collect() - backend_empty_cache(torch_device) - - mixed_int8_config = BitsAndBytesConfig(load_in_8bit=True) - self.model_8bit = SD3Transformer2DModel.from_pretrained( - self.model_name, subfolder="transformer", quantization_config=mixed_int8_config, device_map=torch_device - ) - yield - - def test_training(self): - # Step 1: freeze all parameters - for param in self.model_8bit.parameters(): - param.requires_grad = False # freeze the model - train adapters later - if param.ndim == 1: - # cast the small parameters (e.g. layernorm) to fp32 for stability - param.data = param.data.to(torch.float32) - - # Step 2: add adapters - for _, module in self.model_8bit.named_modules(): - if "Attention" in repr(type(module)): - module.to_k = LoRALayer(module.to_k, rank=4) - module.to_q = LoRALayer(module.to_q, rank=4) - module.to_v = LoRALayer(module.to_v, rank=4) - - # Step 3: dummy batch - input_dict_for_transformer = self.get_dummy_inputs() - model_inputs = { - k: v.to(device=torch_device) for k, v in input_dict_for_transformer.items() if not isinstance(v, bool) - } - model_inputs.update({k: v for k, v in input_dict_for_transformer.items() if k not in model_inputs}) - - # Step 4: Check if the gradient is not None - with torch.amp.autocast(torch_device, dtype=torch.float16): - out = self.model_8bit(**model_inputs)[0] - out.norm().backward() - - for module in self.model_8bit.modules(): - if isinstance(module, LoRALayer): - assert module.adapter[1].weight.grad is not None - assert module.adapter[1].weight.grad.norm().item() > 0 - - -@require_transformers_version_greater("4.44.0") -class TestSlowBnb8bit(Base8bitTests): - @pytest.fixture(autouse=True) - def _setup_slow(self): - gc.collect() - backend_empty_cache(torch_device) - - mixed_int8_config = BitsAndBytesConfig(load_in_8bit=True) - model_8bit = SD3Transformer2DModel.from_pretrained( - self.model_name, subfolder="transformer", quantization_config=mixed_int8_config, device_map=torch_device - ) - self.pipeline_8bit = DiffusionPipeline.from_pretrained( - self.model_name, transformer=model_8bit, torch_dtype=torch.float16 - ) - self.pipeline_8bit.enable_model_cpu_offload() - yield - del self.pipeline_8bit - - gc.collect() - backend_empty_cache(torch_device) - - def test_quality(self): - output = self.pipeline_8bit( - prompt=self.prompt, - num_inference_steps=self.num_inference_steps, - generator=torch.manual_seed(self.seed), - output_type="np", - ).images - out_slice = output[0, -3:, -3:, -1].flatten() - expected_slice = np.array([0.0674, 0.0623, 0.0364, 0.0632, 0.0671, 0.0430, 0.0317, 0.0493, 0.0583]) - - max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice) - assert max_diff < 1e-2 - - def test_model_cpu_offload_raises_warning(self): - model_8bit = SD3Transformer2DModel.from_pretrained( - self.model_name, - subfolder="transformer", - quantization_config=BitsAndBytesConfig(load_in_8bit=True), - device_map=torch_device, - ) - pipeline_8bit = DiffusionPipeline.from_pretrained( - self.model_name, transformer=model_8bit, torch_dtype=torch.float16 - ) - logger = logging.get_logger("diffusers.pipelines.pipeline_utils") - logger.setLevel(30) - - with CaptureLogger(logger) as cap_logger: - pipeline_8bit.enable_model_cpu_offload() - - assert "has been loaded in `bitsandbytes` 8bit" in cap_logger.out - - def test_moving_to_cpu_throws_warning(self): - model_8bit = SD3Transformer2DModel.from_pretrained( - self.model_name, - subfolder="transformer", - quantization_config=BitsAndBytesConfig(load_in_8bit=True), - device_map=torch_device, - ) - logger = logging.get_logger("diffusers.pipelines.pipeline_utils") - logger.setLevel(30) - - with CaptureLogger(logger) as cap_logger: - # Because `model.dtype` will return torch.float16 as SD3 transformer has - # a conv layer as the first layer. - _ = DiffusionPipeline.from_pretrained( - self.model_name, transformer=model_8bit, torch_dtype=torch.float16 - ).to("cpu") - - assert "Pipelines loaded with `dtype=torch.float16`" in cap_logger.out - - def test_generate_quality_dequantize(self): - r""" - Test that loading the model and unquantize it produce correct results. - """ - self.pipeline_8bit.transformer.dequantize() - output = self.pipeline_8bit( - prompt=self.prompt, - num_inference_steps=self.num_inference_steps, - generator=torch.manual_seed(self.seed), - output_type="np", - ).images - - out_slice = output[0, -3:, -3:, -1].flatten() - expected_slice = np.array([0.0266, 0.0264, 0.0271, 0.0110, 0.0310, 0.0098, 0.0078, 0.0256, 0.0208]) - max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice) - assert max_diff < 1e-2 - - # 8bit models cannot be offloaded to CPU. - assert self.pipeline_8bit.transformer.device.type == torch_device - # calling it again shouldn't be a problem - _ = self.pipeline_8bit( - prompt=self.prompt, - num_inference_steps=2, - generator=torch.manual_seed(self.seed), - output_type="np", - ).images - - @pytest.mark.xfail( - condition=is_accelerate_version("<=", "1.1.1"), - reason="Test will pass after https://github.com/huggingface/accelerate/pull/3223 is in a release.", - strict=True, - ) - def test_pipeline_cuda_placement_works_with_mixed_int8(self): - transformer_8bit_config = BitsAndBytesConfig(load_in_8bit=True) - transformer_8bit = SD3Transformer2DModel.from_pretrained( - self.model_name, - subfolder="transformer", - quantization_config=transformer_8bit_config, - torch_dtype=torch.float16, - device_map=torch_device, - ) - text_encoder_3_8bit_config = BnbConfig(load_in_8bit=True) - text_encoder_3_8bit = T5EncoderModel.from_pretrained( - self.model_name, - subfolder="text_encoder_3", - quantization_config=text_encoder_3_8bit_config, - torch_dtype=torch.float16, - device_map=torch_device, - ) - - # CUDA device placement works. - device = torch_device if torch_device != "rocm" else "cuda" - pipeline_8bit = DiffusionPipeline.from_pretrained( - self.model_name, - transformer=transformer_8bit, - text_encoder_3=text_encoder_3_8bit, - torch_dtype=torch.float16, - ).to(device) - - # Check if inference works. - _ = pipeline_8bit(self.prompt, max_sequence_length=20, num_inference_steps=2) - - del pipeline_8bit - - def test_device_map(self): - """ - Test if the quantized model is working properly with "auto" - pu/disk offloading doesn't work with bnb. - """ - - def get_dummy_tensor_inputs(device=None, seed: int = 0): - batch_size = 1 - num_latent_channels = 4 - num_image_channels = 3 - height = width = 4 - sequence_length = 48 - embedding_dim = 32 - - torch.manual_seed(seed) - hidden_states = torch.randn((batch_size, height * width, num_latent_channels)).to( - device, dtype=torch.bfloat16 - ) - - torch.manual_seed(seed) - encoder_hidden_states = torch.randn((batch_size, sequence_length, embedding_dim)).to( - device, dtype=torch.bfloat16 - ) - - torch.manual_seed(seed) - pooled_prompt_embeds = torch.randn((batch_size, embedding_dim)).to(device, dtype=torch.bfloat16) - - torch.manual_seed(seed) - text_ids = torch.randn((sequence_length, num_image_channels)).to(device, dtype=torch.bfloat16) - - torch.manual_seed(seed) - image_ids = torch.randn((height * width, num_image_channels)).to(device, dtype=torch.bfloat16) - - timestep = torch.tensor([1.0]).to(device, dtype=torch.bfloat16).expand(batch_size) - - return { - "hidden_states": hidden_states, - "encoder_hidden_states": encoder_hidden_states, - "pooled_projections": pooled_prompt_embeds, - "txt_ids": text_ids, - "img_ids": image_ids, - "timestep": timestep, - } - - inputs = get_dummy_tensor_inputs(torch_device) - expected_slice = np.array( - [ - 0.33789062, - -0.04736328, - -0.00256348, - -0.23144531, - -0.49804688, - 0.4375, - -0.15429688, - -0.65234375, - 0.44335938, - ] - ) - - # non sharded - quantization_config = BitsAndBytesConfig(load_in_8bit=True) - quantized_model = FluxTransformer2DModel.from_pretrained( - "hf-internal-testing/tiny-flux-pipe", - subfolder="transformer", - quantization_config=quantization_config, - device_map="auto", - torch_dtype=torch.bfloat16, - ) - - weight = quantized_model.transformer_blocks[0].ff.net[2].weight - assert isinstance(weight, bnb.nn.modules.Int8Params) - - output = quantized_model(**inputs)[0] - output_slice = output.flatten()[-9:].detach().float().cpu().numpy() - assert numpy_cosine_similarity_distance(output_slice, expected_slice) < 1e-3 - - # sharded - quantization_config = BitsAndBytesConfig(load_in_8bit=True) - quantized_model = FluxTransformer2DModel.from_pretrained( - "hf-internal-testing/tiny-flux-sharded", - subfolder="transformer", - quantization_config=quantization_config, - device_map="auto", - torch_dtype=torch.bfloat16, - ) - - weight = quantized_model.transformer_blocks[0].ff.net[2].weight - assert isinstance(weight, bnb.nn.modules.Int8Params) - output = quantized_model(**inputs)[0] - output_slice = output.flatten()[-9:].detach().float().cpu().numpy() - - assert numpy_cosine_similarity_distance(output_slice, expected_slice) < 1e-3 - - -@require_transformers_version_greater("4.44.0") -@require_big_accelerator -class TestSlowBnb8bitFlux(Base8bitTests): - @pytest.fixture(autouse=True) - def _setup_slow_flux(self): - gc.collect() - backend_empty_cache(torch_device) - - model_id = "hf-internal-testing/flux.1-dev-int8-pkg" - t5_8bit = T5EncoderModel.from_pretrained(model_id, subfolder="text_encoder_2") - transformer_8bit = FluxTransformer2DModel.from_pretrained(model_id, subfolder="transformer") - self.pipeline_8bit = DiffusionPipeline.from_pretrained( - "black-forest-labs/FLUX.1-dev", - text_encoder_2=t5_8bit, - transformer=transformer_8bit, - torch_dtype=torch.float16, - ) - self.pipeline_8bit.enable_model_cpu_offload() - yield - del self.pipeline_8bit - - gc.collect() - backend_empty_cache(torch_device) - - def test_quality(self): - # keep the resolution and max tokens to a lower number for faster execution. - output = self.pipeline_8bit( - prompt=self.prompt, - num_inference_steps=self.num_inference_steps, - generator=torch.manual_seed(self.seed), - height=256, - width=256, - max_sequence_length=64, - output_type="np", - ).images - out_slice = output[0, -3:, -3:, -1].flatten() - expected_slice = np.array([0.0574, 0.0554, 0.0581, 0.0686, 0.0676, 0.0759, 0.0757, 0.0803, 0.0930]) - - max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice) - assert max_diff < 1e-3 - - @require_peft_version_greater("0.14.0") - def test_lora_loading(self): - self.pipeline_8bit.load_lora_weights( - hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors"), adapter_name="hyper-sd" - ) - self.pipeline_8bit.set_adapters("hyper-sd", adapter_weights=0.125) - - output = self.pipeline_8bit( - prompt=self.prompt, - height=256, - width=256, - max_sequence_length=64, - output_type="np", - num_inference_steps=8, - generator=torch.manual_seed(42), - ).images - out_slice = output[0, -3:, -3:, -1].flatten() - - expected_slice = np.array([0.3916, 0.3916, 0.3887, 0.4243, 0.4155, 0.4233, 0.4570, 0.4531, 0.4248]) - - max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice) - assert max_diff < 1e-3 - - -@require_transformers_version_greater("4.44.0") -@require_peft_backend -class TestSlowBnb4BitFluxControlWithLora(Base8bitTests): - @pytest.fixture(autouse=True) - def _setup_flux_control_lora(self): - gc.collect() - backend_empty_cache(torch_device) - - self.pipeline_8bit = FluxControlPipeline.from_pretrained( - "black-forest-labs/FLUX.1-dev", - quantization_config=PipelineQuantizationConfig( - quant_backend="bitsandbytes_8bit", - quant_kwargs={"load_in_8bit": True}, - components_to_quantize=["transformer", "text_encoder_2"], - ), - torch_dtype=torch.float16, - ) - self.pipeline_8bit.enable_model_cpu_offload() - yield - del self.pipeline_8bit - - gc.collect() - backend_empty_cache(torch_device) - - def test_lora_loading(self): - self.pipeline_8bit.load_lora_weights("black-forest-labs/FLUX.1-Canny-dev-lora") - - output = self.pipeline_8bit( - prompt=self.prompt, - control_image=Image.new(mode="RGB", size=(256, 256)), - height=256, - width=256, - max_sequence_length=64, - output_type="np", - num_inference_steps=8, - generator=torch.Generator().manual_seed(42), - ).images - out_slice = output[0, -3:, -3:, -1].flatten() - # Hardware-dependent: the Control LoRA dequantizes and expands `x_embedder`, and the error - # accumulates over the 8 denoising steps enough that even different CUDA GPUs disagree, so - # reference slices are stored per accelerator backend. - expected_slices = Expectations( - { - (None, None): np.array([0.2029, 0.2136, 0.2268, 0.1921, 0.1997, 0.2185, 0.2021, 0.2183, 0.2292]), - ("xpu", 5): np.array([0.0955, 0.1223, 0.1509, 0.0872, 0.1155, 0.1890, 0.0754, 0.1028, 0.2178]), - } - ) - expected_slice = expected_slices.get_expectation() - - max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice) - assert max_diff < 1e-3, f"{out_slice=} != {expected_slice=}" - - -@slow -class TestBnb8bitSerialization(Base8bitTests): - @pytest.fixture(autouse=True) - def _setup_serialization(self): - gc.collect() - backend_empty_cache(torch_device) - - quantization_config = BitsAndBytesConfig( - load_in_8bit=True, - ) - self.model_0 = SD3Transformer2DModel.from_pretrained( - self.model_name, subfolder="transformer", quantization_config=quantization_config, device_map=torch_device - ) - yield - del self.model_0 - - gc.collect() - backend_empty_cache(torch_device) - - def test_serialization(self): - r""" - Test whether it is possible to serialize a model in 8-bit. Uses most typical params as default. - """ - assert "_pre_quantization_dtype" in self.model_0.config - with tempfile.TemporaryDirectory() as tmpdirname: - self.model_0.save_pretrained(tmpdirname) - - config = SD3Transformer2DModel.load_config(tmpdirname) - assert "quantization_config" in config - assert "_pre_quantization_dtype" not in config - - model_1 = SD3Transformer2DModel.from_pretrained(tmpdirname) - - # checking quantized linear module weight - linear = get_some_linear_layer(model_1) - assert linear.weight.__class__ == bnb.nn.Int8Params - assert hasattr(linear.weight, "SCB") - - # checking memory footpring - assert self.model_0.get_memory_footprint() / model_1.get_memory_footprint() == pytest.approx(1, abs=10**-2) - - # Matching all parameters and their quant_state items: - d0 = dict(self.model_0.named_parameters()) - d1 = dict(model_1.named_parameters()) - assert d0.keys() == d1.keys() - - # comparing forward() outputs - dummy_inputs = self.get_dummy_inputs() - inputs = {k: v.to(torch_device) for k, v in dummy_inputs.items() if isinstance(v, torch.Tensor)} - inputs.update({k: v for k, v in dummy_inputs.items() if k not in inputs}) - out_0 = self.model_0(**inputs)[0] - out_1 = model_1(**inputs)[0] - assert torch.equal(out_0, out_1) - - def test_serialization_sharded(self): - with tempfile.TemporaryDirectory() as tmpdirname: - self.model_0.save_pretrained(tmpdirname, max_shard_size="200MB") - - config = SD3Transformer2DModel.load_config(tmpdirname) - assert "quantization_config" in config - assert "_pre_quantization_dtype" not in config - - model_1 = SD3Transformer2DModel.from_pretrained(tmpdirname) - - # checking quantized linear module weight - linear = get_some_linear_layer(model_1) - assert linear.weight.__class__ == bnb.nn.Int8Params - assert hasattr(linear.weight, "SCB") - - # comparing forward() outputs - dummy_inputs = self.get_dummy_inputs() - inputs = {k: v.to(torch_device) for k, v in dummy_inputs.items() if isinstance(v, torch.Tensor)} - inputs.update({k: v for k, v in dummy_inputs.items() if k not in inputs}) - out_0 = self.model_0(**inputs)[0] - out_1 = model_1(**inputs)[0] - assert torch.equal(out_0, out_1) - - -@require_torch_version_greater_equal("2.6.0") -@require_bitsandbytes_version_greater("0.48.0") -class TestBnb8BitCompile(QuantCompileTests): - @property - def quantization_config(self): - return PipelineQuantizationConfig( - quant_backend="bitsandbytes_8bit", - quant_kwargs={"load_in_8bit": True}, - components_to_quantize=["transformer", "text_encoder_2"], - ) - - def test_torch_compile_with_cpu_offload(self): - super()._test_torch_compile_with_cpu_offload(torch_dtype=torch.float16) diff --git a/tests/quantization/gguf/test_gguf.py b/tests/quantization/gguf/test_gguf.py index 83c8c5d9ef98..97f3c186f6d0 100644 --- a/tests/quantization/gguf/test_gguf.py +++ b/tests/quantization/gguf/test_gguf.py @@ -1,56 +1,37 @@ import gc -import numpy as np import pytest import torch import torch.nn as nn -from diffusers import ( - AuraFlowPipeline, - AuraFlowTransformer2DModel, - DiffusionPipeline, - FluxControlPipeline, - FluxPipeline, - FluxTransformer2DModel, - GGUFQuantizationConfig, - HiDreamImageTransformer2DModel, - SD3Transformer2DModel, - StableDiffusion3Pipeline, - WanAnimateTransformer3DModel, - WanTransformer3DModel, - WanVACETransformer3DModel, -) -from diffusers.utils import load_image - from ...testing_utils import ( - Expectations, backend_empty_cache, - backend_max_memory_allocated, - backend_reset_peak_memory_stats, enable_full_determinism, + is_gguf, is_gguf_available, + is_quantization, nightly, - numpy_cosine_similarity_distance, require_accelerate, require_accelerator, - require_big_accelerator, require_gguf_version_greater_or_equal, require_kernels_version_greater_or_equal, - require_peft_backend, - require_torch_version_greater, torch_device, ) -from ..test_torch_compile_utils import QuantCompileTests if is_gguf_available(): import gguf - from diffusers.quantizers.gguf.utils import GGUFLinear, GGUFParameter + from diffusers.quantizers.gguf.utils import GGUFParameter enable_full_determinism() +# Model-level GGUF tests live in `tests/models/testing_utils/quantization.py` and pipeline-level +# ones in `tests/pipelines/testing_utils/quantization.py`. This module covers backend behavior that +# fits neither: CUDA kernel correctness. +@is_quantization +@is_gguf @nightly @require_accelerate @require_accelerator @@ -102,667 +83,3 @@ def test_cuda_kernels_vs_native(self): assert torch.allclose(output_native, output_cuda, 1e-2), ( f"GGUF CUDA Kernel Output is different from Native Output for {quant_type}" ) - - -@nightly -@require_big_accelerator -@require_accelerate -@require_gguf_version_greater_or_equal("0.10.0") -class GGUFSingleFileTesterMixin: - ckpt_path = None - model_cls = None - torch_dtype = torch.bfloat16 - expected_memory_use_in_gb = 5 - - def test_gguf_parameters(self): - quant_storage_type = torch.uint8 - quantization_config = GGUFQuantizationConfig(compute_dtype=self.torch_dtype) - model = self.model_cls.from_single_file(self.ckpt_path, quantization_config=quantization_config) - - for param_name, param in model.named_parameters(): - if isinstance(param, GGUFParameter): - assert hasattr(param, "quant_type") - assert param.dtype == quant_storage_type - - def test_gguf_linear_layers(self): - quantization_config = GGUFQuantizationConfig(compute_dtype=self.torch_dtype) - model = self.model_cls.from_single_file(self.ckpt_path, quantization_config=quantization_config) - - for name, module in model.named_modules(): - if isinstance(module, torch.nn.Linear) and hasattr(module.weight, "quant_type"): - assert module.weight.dtype == torch.uint8 - if module.bias is not None: - assert module.bias.dtype == self.torch_dtype - - def test_gguf_memory_usage(self): - quantization_config = GGUFQuantizationConfig(compute_dtype=self.torch_dtype) - - model = self.model_cls.from_single_file( - self.ckpt_path, quantization_config=quantization_config, torch_dtype=self.torch_dtype - ) - model.to(torch_device) - assert (model.get_memory_footprint() / 1024**3) < self.expected_memory_use_in_gb - inputs = self.get_dummy_inputs() - - backend_reset_peak_memory_stats(torch_device) - backend_empty_cache(torch_device) - with torch.no_grad(): - model(**inputs) - max_memory = backend_max_memory_allocated(torch_device) - assert (max_memory / 1024**3) < self.expected_memory_use_in_gb - - def test_keep_modules_in_fp32(self): - r""" - A simple tests to check if the modules under `_keep_in_fp32_modules` are kept in fp32. - Also ensures if inference works. - """ - _keep_in_fp32_modules = self.model_cls._keep_in_fp32_modules - self.model_cls._keep_in_fp32_modules = ["proj_out"] - - quantization_config = GGUFQuantizationConfig(compute_dtype=self.torch_dtype) - model = self.model_cls.from_single_file(self.ckpt_path, quantization_config=quantization_config) - - for name, module in model.named_modules(): - if isinstance(module, torch.nn.Linear): - if name in model._keep_in_fp32_modules: - assert module.weight.dtype == torch.float32 - self.model_cls._keep_in_fp32_modules = _keep_in_fp32_modules - - def test_dtype_assignment(self): - quantization_config = GGUFQuantizationConfig(compute_dtype=self.torch_dtype) - model = self.model_cls.from_single_file(self.ckpt_path, quantization_config=quantization_config) - - with pytest.raises(ValueError): - # Tries with a `dtype` - model.to(torch.float16) - - with pytest.raises(ValueError): - # Tries with a `device` and `dtype` - device_0 = f"{torch_device}:0" - model.to(device=device_0, dtype=torch.float16) - - with pytest.raises(ValueError): - # Tries with a cast - model.float() - - with pytest.raises(ValueError): - # Tries with a cast - model.half() - - # This should work - model.to(torch_device) - - def test_dequantize_model(self): - quantization_config = GGUFQuantizationConfig(compute_dtype=self.torch_dtype) - model = self.model_cls.from_single_file(self.ckpt_path, quantization_config=quantization_config) - model.dequantize() - - def _check_for_gguf_linear(model): - has_children = list(model.children()) - if not has_children: - return - - for name, module in model.named_children(): - if isinstance(module, nn.Linear): - assert not isinstance(module, GGUFLinear), f"{name} is still GGUFLinear" - assert not isinstance(module.weight, GGUFParameter), f"{name} weight is still GGUFParameter" - - for name, module in model.named_children(): - _check_for_gguf_linear(module) - - -class TestFluxGGUFSingleFile(GGUFSingleFileTesterMixin): - ckpt_path = "https://huggingface.co/city96/FLUX.1-dev-gguf/blob/main/flux1-dev-Q2_K.gguf" - diffusers_ckpt_path = "https://huggingface.co/sayakpaul/flux-diffusers-gguf/blob/main/model-Q4_0.gguf" - torch_dtype = torch.bfloat16 - model_cls = FluxTransformer2DModel - expected_memory_use_in_gb = 5 - - @pytest.fixture(autouse=True) - def _setup_flux(self): - gc.collect() - backend_empty_cache(torch_device) - yield - gc.collect() - backend_empty_cache(torch_device) - - def get_dummy_inputs(self): - return { - "hidden_states": torch.randn((1, 4096, 64), generator=torch.Generator("cpu").manual_seed(0)).to( - torch_device, self.torch_dtype - ), - "encoder_hidden_states": torch.randn( - (1, 512, 4096), - generator=torch.Generator("cpu").manual_seed(0), - ).to(torch_device, self.torch_dtype), - "pooled_projections": torch.randn( - (1, 768), - generator=torch.Generator("cpu").manual_seed(0), - ).to(torch_device, self.torch_dtype), - "timestep": torch.tensor([1]).to(torch_device, self.torch_dtype), - "img_ids": torch.randn((4096, 3), generator=torch.Generator("cpu").manual_seed(0)).to( - torch_device, self.torch_dtype - ), - "txt_ids": torch.randn((512, 3), generator=torch.Generator("cpu").manual_seed(0)).to( - torch_device, self.torch_dtype - ), - "guidance": torch.tensor([3.5]).to(torch_device, self.torch_dtype), - } - - def test_pipeline_inference(self): - quantization_config = GGUFQuantizationConfig(compute_dtype=self.torch_dtype) - transformer = self.model_cls.from_single_file( - self.ckpt_path, quantization_config=quantization_config, torch_dtype=self.torch_dtype - ) - pipe = FluxPipeline.from_pretrained( - "black-forest-labs/FLUX.1-dev", transformer=transformer, torch_dtype=self.torch_dtype - ) - pipe.enable_model_cpu_offload() - - prompt = "a cat holding a sign that says hello" - output = pipe( - prompt=prompt, num_inference_steps=2, generator=torch.Generator("cpu").manual_seed(0), output_type="np" - ).images[0] - output_slice = output[:3, :3, :].flatten() - expected_slice = np.array( - [ - 0.47265625, - 0.43359375, - 0.359375, - 0.47070312, - 0.421875, - 0.34375, - 0.46875, - 0.421875, - 0.34765625, - 0.46484375, - 0.421875, - 0.34179688, - 0.47070312, - 0.42578125, - 0.34570312, - 0.46875, - 0.42578125, - 0.3515625, - 0.45507812, - 0.4140625, - 0.33984375, - 0.4609375, - 0.41796875, - 0.34375, - 0.45898438, - 0.41796875, - 0.34375, - ] - ) - max_diff = numpy_cosine_similarity_distance(expected_slice, output_slice) - assert max_diff < 1e-4 - - def test_loading_gguf_diffusers_format(self): - model = self.model_cls.from_single_file( - self.diffusers_ckpt_path, - subfolder="transformer", - quantization_config=GGUFQuantizationConfig(compute_dtype=torch.bfloat16), - config="black-forest-labs/FLUX.1-dev", - ) - model.to(torch_device) - model(**self.get_dummy_inputs()) - - -class TestSD35LargeGGUFSingleFile(GGUFSingleFileTesterMixin): - ckpt_path = "https://huggingface.co/city96/stable-diffusion-3.5-large-gguf/blob/main/sd3.5_large-Q4_0.gguf" - torch_dtype = torch.bfloat16 - model_cls = SD3Transformer2DModel - expected_memory_use_in_gb = 5 - - @pytest.fixture(autouse=True) - def _setup_sd35_large(self): - gc.collect() - backend_empty_cache(torch_device) - yield - gc.collect() - backend_empty_cache(torch_device) - - def get_dummy_inputs(self): - return { - "hidden_states": torch.randn((1, 16, 64, 64), generator=torch.Generator("cpu").manual_seed(0)).to( - torch_device, self.torch_dtype - ), - "encoder_hidden_states": torch.randn( - (1, 512, 4096), - generator=torch.Generator("cpu").manual_seed(0), - ).to(torch_device, self.torch_dtype), - "pooled_projections": torch.randn( - (1, 2048), - generator=torch.Generator("cpu").manual_seed(0), - ).to(torch_device, self.torch_dtype), - "timestep": torch.tensor([1]).to(torch_device, self.torch_dtype), - } - - def test_pipeline_inference(self): - quantization_config = GGUFQuantizationConfig(compute_dtype=self.torch_dtype) - transformer = self.model_cls.from_single_file( - self.ckpt_path, quantization_config=quantization_config, torch_dtype=self.torch_dtype - ) - pipe = StableDiffusion3Pipeline.from_pretrained( - "stabilityai/stable-diffusion-3.5-large", transformer=transformer, torch_dtype=self.torch_dtype - ) - pipe.enable_model_cpu_offload() - - prompt = "a cat holding a sign that says hello" - output = pipe( - prompt=prompt, - num_inference_steps=2, - generator=torch.Generator("cpu").manual_seed(0), - output_type="np", - ).images[0] - output_slice = output[:3, :3, :].flatten() - expected_slices = Expectations( - { - ("xpu", 3): np.array( - [ - 0.16796875, - 0.27929688, - 0.28320312, - 0.11328125, - 0.27539062, - 0.26171875, - 0.10742188, - 0.26367188, - 0.26171875, - 0.1484375, - 0.2734375, - 0.296875, - 0.13476562, - 0.2890625, - 0.30078125, - 0.1171875, - 0.28125, - 0.28125, - 0.16015625, - 0.31445312, - 0.30078125, - 0.15625, - 0.32421875, - 0.296875, - 0.14453125, - 0.30859375, - 0.2890625, - ] - ), - ("cuda", 7): np.array( - [ - 0.17578125, - 0.27539062, - 0.27734375, - 0.11914062, - 0.26953125, - 0.25390625, - 0.109375, - 0.25390625, - 0.25, - 0.15039062, - 0.26171875, - 0.28515625, - 0.13671875, - 0.27734375, - 0.28515625, - 0.12109375, - 0.26757812, - 0.265625, - 0.16210938, - 0.29882812, - 0.28515625, - 0.15625, - 0.30664062, - 0.27734375, - 0.14648438, - 0.29296875, - 0.26953125, - ] - ), - } - ) - expected_slice = expected_slices.get_expectation() - max_diff = numpy_cosine_similarity_distance(expected_slice, output_slice) - assert max_diff < 1e-4 - - -class TestSD35MediumGGUFSingleFile(GGUFSingleFileTesterMixin): - ckpt_path = "https://huggingface.co/city96/stable-diffusion-3.5-medium-gguf/blob/main/sd3.5_medium-Q3_K_M.gguf" - torch_dtype = torch.bfloat16 - model_cls = SD3Transformer2DModel - expected_memory_use_in_gb = 2 - - @pytest.fixture(autouse=True) - def _setup_sd35_medium(self): - gc.collect() - backend_empty_cache(torch_device) - yield - gc.collect() - backend_empty_cache(torch_device) - - def get_dummy_inputs(self): - return { - "hidden_states": torch.randn((1, 16, 64, 64), generator=torch.Generator("cpu").manual_seed(0)).to( - torch_device, self.torch_dtype - ), - "encoder_hidden_states": torch.randn( - (1, 512, 4096), - generator=torch.Generator("cpu").manual_seed(0), - ).to(torch_device, self.torch_dtype), - "pooled_projections": torch.randn( - (1, 2048), - generator=torch.Generator("cpu").manual_seed(0), - ).to(torch_device, self.torch_dtype), - "timestep": torch.tensor([1]).to(torch_device, self.torch_dtype), - } - - def test_pipeline_inference(self): - quantization_config = GGUFQuantizationConfig(compute_dtype=self.torch_dtype) - transformer = self.model_cls.from_single_file( - self.ckpt_path, quantization_config=quantization_config, torch_dtype=self.torch_dtype - ) - pipe = StableDiffusion3Pipeline.from_pretrained( - "stabilityai/stable-diffusion-3.5-medium", transformer=transformer, torch_dtype=self.torch_dtype - ) - pipe.enable_model_cpu_offload() - - prompt = "a cat holding a sign that says hello" - output = pipe( - prompt=prompt, num_inference_steps=2, generator=torch.Generator("cpu").manual_seed(0), output_type="np" - ).images[0] - output_slice = output[:3, :3, :].flatten() - expected_slice = np.array( - [ - 0.625, - 0.6171875, - 0.609375, - 0.65625, - 0.65234375, - 0.640625, - 0.6484375, - 0.640625, - 0.625, - 0.6484375, - 0.63671875, - 0.6484375, - 0.66796875, - 0.65625, - 0.65234375, - 0.6640625, - 0.6484375, - 0.6328125, - 0.6640625, - 0.6484375, - 0.640625, - 0.67578125, - 0.66015625, - 0.62109375, - 0.671875, - 0.65625, - 0.62109375, - ] - ) - max_diff = numpy_cosine_similarity_distance(expected_slice, output_slice) - assert max_diff < 1e-4 - - -class TestAuraFlowGGUFSingleFile(GGUFSingleFileTesterMixin): - ckpt_path = "https://huggingface.co/city96/AuraFlow-v0.3-gguf/blob/main/aura_flow_0.3-Q2_K.gguf" - torch_dtype = torch.bfloat16 - model_cls = AuraFlowTransformer2DModel - expected_memory_use_in_gb = 4 - - @pytest.fixture(autouse=True) - def _setup_auraflow(self): - gc.collect() - backend_empty_cache(torch_device) - yield - gc.collect() - backend_empty_cache(torch_device) - - def get_dummy_inputs(self): - return { - "hidden_states": torch.randn((1, 4, 64, 64), generator=torch.Generator("cpu").manual_seed(0)).to( - torch_device, self.torch_dtype - ), - "encoder_hidden_states": torch.randn( - (1, 512, 2048), - generator=torch.Generator("cpu").manual_seed(0), - ).to(torch_device, self.torch_dtype), - "timestep": torch.tensor([1]).to(torch_device, self.torch_dtype), - } - - def test_pipeline_inference(self): - quantization_config = GGUFQuantizationConfig(compute_dtype=self.torch_dtype) - transformer = self.model_cls.from_single_file( - self.ckpt_path, quantization_config=quantization_config, torch_dtype=self.torch_dtype - ) - pipe = AuraFlowPipeline.from_pretrained( - "fal/AuraFlow-v0.3", transformer=transformer, torch_dtype=self.torch_dtype - ) - pipe.enable_model_cpu_offload() - - prompt = "a pony holding a sign that says hello" - output = pipe( - prompt=prompt, num_inference_steps=2, generator=torch.Generator("cpu").manual_seed(0), output_type="np" - ).images[0] - output_slice = output[:3, :3, :].flatten() - expected_slice = np.array( - [ - 0.46484375, - 0.546875, - 0.64453125, - 0.48242188, - 0.53515625, - 0.59765625, - 0.47070312, - 0.5078125, - 0.5703125, - 0.42773438, - 0.50390625, - 0.5703125, - 0.47070312, - 0.515625, - 0.57421875, - 0.45898438, - 0.48632812, - 0.53515625, - 0.4453125, - 0.5078125, - 0.56640625, - 0.47851562, - 0.5234375, - 0.57421875, - 0.48632812, - 0.5234375, - 0.56640625, - ] - ) - max_diff = numpy_cosine_similarity_distance(expected_slice, output_slice) - assert max_diff < 1e-4 - - -@require_peft_backend -@nightly -@require_big_accelerator -@require_accelerate -@require_gguf_version_greater_or_equal("0.10.0") -class TestFluxControlLoRAGGUF: - def test_lora_loading(self): - ckpt_path = "https://huggingface.co/city96/FLUX.1-dev-gguf/blob/main/flux1-dev-Q2_K.gguf" - transformer = FluxTransformer2DModel.from_single_file( - ckpt_path, - quantization_config=GGUFQuantizationConfig(compute_dtype=torch.bfloat16), - torch_dtype=torch.bfloat16, - ) - pipe = FluxControlPipeline.from_pretrained( - "black-forest-labs/FLUX.1-dev", - transformer=transformer, - torch_dtype=torch.bfloat16, - ).to(torch_device) - pipe.load_lora_weights("black-forest-labs/FLUX.1-Canny-dev-lora") - - prompt = "A robot made of exotic candies and chocolates of different kinds. The background is filled with confetti and celebratory gifts." - control_image = load_image( - "https://huggingface.co/datasets/sayakpaul/sample-datasets/resolve/main/control_image_robot_canny.png" - ) - - output = pipe( - prompt=prompt, - control_image=control_image, - height=256, - width=256, - num_inference_steps=10, - guidance_scale=30.0, - output_type="np", - generator=torch.manual_seed(0), - ).images - - out_slice = output[0, -3:, -3:, -1].flatten() - expected_slice = np.array([0.8047, 0.8359, 0.8711, 0.6875, 0.7070, 0.7383, 0.5469, 0.5820, 0.6641]) - - max_diff = numpy_cosine_similarity_distance(expected_slice, out_slice) - assert max_diff < 1e-3 - - -class TestHiDreamGGUFSingleFile(GGUFSingleFileTesterMixin): - ckpt_path = "https://huggingface.co/city96/HiDream-I1-Dev-gguf/blob/main/hidream-i1-dev-Q2_K.gguf" - torch_dtype = torch.bfloat16 - model_cls = HiDreamImageTransformer2DModel - expected_memory_use_in_gb = 8 - - def get_dummy_inputs(self): - return { - "hidden_states": torch.randn((1, 16, 128, 128), generator=torch.Generator("cpu").manual_seed(0)).to( - torch_device, self.torch_dtype - ), - "encoder_hidden_states_t5": torch.randn( - (1, 128, 4096), - generator=torch.Generator("cpu").manual_seed(0), - ).to(torch_device, self.torch_dtype), - "encoder_hidden_states_llama3": torch.randn( - (32, 1, 128, 4096), - generator=torch.Generator("cpu").manual_seed(0), - ).to(torch_device, self.torch_dtype), - "pooled_embeds": torch.randn( - (1, 2048), - generator=torch.Generator("cpu").manual_seed(0), - ).to(torch_device, self.torch_dtype), - "timesteps": torch.tensor([1]).to(torch_device, self.torch_dtype), - } - - -class TestWanGGUFTexttoVideoSingleFile(GGUFSingleFileTesterMixin): - ckpt_path = "https://huggingface.co/city96/Wan2.1-T2V-14B-gguf/blob/main/wan2.1-t2v-14b-Q3_K_S.gguf" - torch_dtype = torch.bfloat16 - model_cls = WanTransformer3DModel - expected_memory_use_in_gb = 9 - - def get_dummy_inputs(self): - return { - "hidden_states": torch.randn((1, 16, 2, 64, 64), generator=torch.Generator("cpu").manual_seed(0)).to( - torch_device, self.torch_dtype - ), - "encoder_hidden_states": torch.randn( - (1, 512, 4096), - generator=torch.Generator("cpu").manual_seed(0), - ).to(torch_device, self.torch_dtype), - "timestep": torch.tensor([1]).to(torch_device, self.torch_dtype), - } - - -class TestWanGGUFImagetoVideoSingleFile(GGUFSingleFileTesterMixin): - ckpt_path = "https://huggingface.co/city96/Wan2.1-I2V-14B-480P-gguf/blob/main/wan2.1-i2v-14b-480p-Q3_K_S.gguf" - torch_dtype = torch.bfloat16 - model_cls = WanTransformer3DModel - expected_memory_use_in_gb = 9 - - def get_dummy_inputs(self): - return { - "hidden_states": torch.randn((1, 36, 2, 64, 64), generator=torch.Generator("cpu").manual_seed(0)).to( - torch_device, self.torch_dtype - ), - "encoder_hidden_states": torch.randn( - (1, 512, 4096), - generator=torch.Generator("cpu").manual_seed(0), - ).to(torch_device, self.torch_dtype), - "encoder_hidden_states_image": torch.randn( - (1, 257, 1280), generator=torch.Generator("cpu").manual_seed(0) - ).to(torch_device, self.torch_dtype), - "timestep": torch.tensor([1]).to(torch_device, self.torch_dtype), - } - - -class TestWanVACEGGUFSingleFile(GGUFSingleFileTesterMixin): - ckpt_path = "https://huggingface.co/QuantStack/Wan2.1_14B_VACE-GGUF/blob/main/Wan2.1_14B_VACE-Q3_K_S.gguf" - torch_dtype = torch.bfloat16 - model_cls = WanVACETransformer3DModel - expected_memory_use_in_gb = 9 - - def get_dummy_inputs(self): - return { - "hidden_states": torch.randn((1, 16, 2, 64, 64), generator=torch.Generator("cpu").manual_seed(0)).to( - torch_device, self.torch_dtype - ), - "encoder_hidden_states": torch.randn( - (1, 512, 4096), - generator=torch.Generator("cpu").manual_seed(0), - ).to(torch_device, self.torch_dtype), - "control_hidden_states": torch.randn( - (1, 96, 2, 64, 64), - generator=torch.Generator("cpu").manual_seed(0), - ).to(torch_device, self.torch_dtype), - "control_hidden_states_scale": torch.randn( - (8,), - generator=torch.Generator("cpu").manual_seed(0), - ).to(torch_device, self.torch_dtype), - "timestep": torch.tensor([1]).to(torch_device, self.torch_dtype), - } - - -class TestWanAnimateGGUFSingleFile(GGUFSingleFileTesterMixin): - ckpt_path = "https://huggingface.co/QuantStack/Wan2.2-Animate-14B-GGUF/blob/main/Wan2.2-Animate-14B-Q3_K_S.gguf" - torch_dtype = torch.bfloat16 - model_cls = WanAnimateTransformer3DModel - expected_memory_use_in_gb = 9 - - def get_dummy_inputs(self): - return { - "hidden_states": torch.randn((1, 16, 2, 64, 64), generator=torch.Generator("cpu").manual_seed(0)).to( - torch_device, self.torch_dtype - ), - "encoder_hidden_states": torch.randn( - (1, 512, 4096), - generator=torch.Generator("cpu").manual_seed(0), - ).to(torch_device, self.torch_dtype), - "control_hidden_states": torch.randn( - (1, 96, 2, 64, 64), - generator=torch.Generator("cpu").manual_seed(0), - ).to(torch_device, self.torch_dtype), - "control_hidden_states_scale": torch.randn( - (8,), - generator=torch.Generator("cpu").manual_seed(0), - ).to(torch_device, self.torch_dtype), - "timestep": torch.tensor([1]).to(torch_device, self.torch_dtype), - } - - -@require_torch_version_greater("2.7.1") -class TestGGUFCompile(QuantCompileTests): - torch_dtype = torch.bfloat16 - gguf_ckpt = "https://huggingface.co/city96/FLUX.1-dev-gguf/blob/main/flux1-dev-Q2_K.gguf" - - @property - def quantization_config(self): - return GGUFQuantizationConfig(compute_dtype=self.torch_dtype) - - def _init_pipeline(self, *args, **kwargs): - transformer = FluxTransformer2DModel.from_single_file( - self.gguf_ckpt, quantization_config=self.quantization_config, torch_dtype=self.torch_dtype - ) - pipe = DiffusionPipeline.from_pretrained( - "black-forest-labs/FLUX.1-dev", transformer=transformer, torch_dtype=self.torch_dtype - ) - return pipe diff --git a/tests/quantization/modelopt/__init__.py b/tests/quantization/modelopt/__init__.py deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/tests/quantization/modelopt/test_modelopt.py b/tests/quantization/modelopt/test_modelopt.py deleted file mode 100644 index 8f104c0bb898..000000000000 --- a/tests/quantization/modelopt/test_modelopt.py +++ /dev/null @@ -1,343 +0,0 @@ -import copy -import gc -import os -import tempfile - -import pytest - -from diffusers import NVIDIAModelOptConfig, SD3Transformer2DModel, StableDiffusion3Pipeline -from diffusers.utils import is_nvidia_modelopt_available, is_torch_available - -from ...testing_utils import ( - backend_empty_cache, - backend_reset_peak_memory_stats, - enable_full_determinism, - nightly, - numpy_cosine_similarity_distance, - require_accelerate, - require_big_accelerator, - require_modelopt_version_greater_or_equal, - require_torch_cuda_compatibility, - torch_device, -) - - -if is_nvidia_modelopt_available(): - import modelopt.torch.opt as mto - import modelopt.torch.quantization as mtq - -if is_torch_available(): - import torch - - from ..utils import LoRALayer, get_memory_consumption_stat - -enable_full_determinism() - - -@nightly -@require_big_accelerator -@require_accelerate -@require_modelopt_version_greater_or_equal("0.33.1") -class ModelOptBaseTesterMixin: - model_id = "hf-internal-testing/tiny-sd3-pipe" - model_cls = SD3Transformer2DModel - pipeline_cls = StableDiffusion3Pipeline - torch_dtype = torch.bfloat16 - expected_memory_reduction = 0.0 - keep_in_fp32_module = "" - modules_to_not_convert = "" - _test_torch_compile = False - - @pytest.fixture(autouse=True) - def _setup(self): - backend_reset_peak_memory_stats(torch_device) - backend_empty_cache(torch_device) - gc.collect() - yield - backend_reset_peak_memory_stats(torch_device) - backend_empty_cache(torch_device) - gc.collect() - - def get_dummy_init_kwargs(self): - return {"quant_type": "FP8"} - - def get_dummy_model_init_kwargs(self): - return { - "pretrained_model_name_or_path": self.model_id, - "torch_dtype": self.torch_dtype, - "quantization_config": NVIDIAModelOptConfig(**self.get_dummy_init_kwargs()), - "subfolder": "transformer", - } - - def test_modelopt_layers(self): - model = self.model_cls.from_pretrained(**self.get_dummy_model_init_kwargs()) - for name, module in model.named_modules(): - if isinstance(module, torch.nn.Linear): - assert mtq.utils.is_quantized(module) - - def test_modelopt_memory_usage(self): - inputs = self.get_dummy_inputs() - inputs = { - k: v.to(device=torch_device, dtype=torch.bfloat16) for k, v in inputs.items() if not isinstance(v, bool) - } - - unquantized_model = self.model_cls.from_pretrained( - self.model_id, torch_dtype=self.torch_dtype, subfolder="transformer" - ) - unquantized_model.to(torch_device) - unquantized_model_memory = get_memory_consumption_stat(unquantized_model, inputs) - - quantized_model = self.model_cls.from_pretrained(**self.get_dummy_model_init_kwargs()) - quantized_model.to(torch_device) - quantized_model_memory = get_memory_consumption_stat(quantized_model, inputs) - - assert unquantized_model_memory / quantized_model_memory >= self.expected_memory_reduction - - def test_keep_modules_in_fp32(self): - _keep_in_fp32_modules = self.model_cls._keep_in_fp32_modules - self.model_cls._keep_in_fp32_modules = self.keep_in_fp32_module - - model = self.model_cls.from_pretrained(**self.get_dummy_model_init_kwargs()) - model.to(torch_device) - - for name, module in model.named_modules(): - if isinstance(module, torch.nn.Linear): - if name in model._keep_in_fp32_modules: - assert module.weight.dtype == torch.float32 - self.model_cls._keep_in_fp32_modules = _keep_in_fp32_modules - - def test_modules_to_not_convert(self): - init_kwargs = self.get_dummy_model_init_kwargs() - quantization_config_kwargs = self.get_dummy_init_kwargs() - quantization_config_kwargs.update({"modules_to_not_convert": self.modules_to_not_convert}) - quantization_config = NVIDIAModelOptConfig(**quantization_config_kwargs) - init_kwargs.update({"quantization_config": quantization_config}) - - model = self.model_cls.from_pretrained(**init_kwargs) - model.to(torch_device) - - for name, module in model.named_modules(): - if name in self.modules_to_not_convert: - assert not mtq.utils.is_quantized(module) - - def test_dtype_assignment(self): - model = self.model_cls.from_pretrained(**self.get_dummy_model_init_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) - - def test_serialization(self): - model = self.model_cls.from_pretrained(**self.get_dummy_model_init_kwargs()) - inputs = self.get_dummy_inputs() - - model.to(torch_device) - with torch.no_grad(): - model_output = model(**inputs) - - with tempfile.TemporaryDirectory() as tmp_dir: - model.save_pretrained(tmp_dir) - saved_model = self.model_cls.from_pretrained( - tmp_dir, - torch_dtype=torch.bfloat16, - ) - - saved_model.to(torch_device) - with torch.no_grad(): - saved_model_output = saved_model(**inputs) - - assert torch.allclose(model_output.sample, saved_model_output.sample, rtol=1e-5, atol=1e-5) - - def test_torch_compile(self): - if not self._test_torch_compile: - return - - model = self.model_cls.from_pretrained(**self.get_dummy_model_init_kwargs()) - compiled_model = torch.compile(model, mode="max-autotune", fullgraph=True, dynamic=False) - - model.to(torch_device) - with torch.no_grad(): - model_output = model(**self.get_dummy_inputs()).sample - - compiled_model.to(torch_device) - with torch.no_grad(): - compiled_model_output = compiled_model(**self.get_dummy_inputs()).sample - - model_output = model_output.detach().float().cpu().numpy() - compiled_model_output = compiled_model_output.detach().float().cpu().numpy() - - max_diff = numpy_cosine_similarity_distance(model_output.flatten(), compiled_model_output.flatten()) - assert max_diff < 1e-3 - - def test_device_map_error(self): - with pytest.raises(ValueError): - _ = self.model_cls.from_pretrained( - **self.get_dummy_model_init_kwargs(), - device_map={0: "8GB", "cpu": "16GB"}, - ) - - def get_dummy_inputs(self): - batch_size = 1 - seq_len = 16 - height = width = 32 - num_latent_channels = 4 - caption_channels = 8 - - torch.manual_seed(0) - hidden_states = torch.randn((batch_size, num_latent_channels, height, width)).to( - torch_device, dtype=torch.bfloat16 - ) - encoder_hidden_states = torch.randn((batch_size, seq_len, caption_channels)).to( - torch_device, dtype=torch.bfloat16 - ) - timestep = torch.tensor([1.0]).to(torch_device, dtype=torch.bfloat16).expand(batch_size) - - return { - "hidden_states": hidden_states, - "encoder_hidden_states": encoder_hidden_states, - "timestep": timestep, - } - - def test_model_cpu_offload(self): - init_kwargs = self.get_dummy_init_kwargs() - transformer = self.model_cls.from_pretrained( - self.model_id, - quantization_config=NVIDIAModelOptConfig(**init_kwargs), - subfolder="transformer", - torch_dtype=torch.bfloat16, - ) - pipe = self.pipeline_cls.from_pretrained(self.model_id, transformer=transformer, torch_dtype=torch.bfloat16) - pipe.enable_model_cpu_offload(device=torch_device) - _ = pipe("a cat holding a sign that says hello", num_inference_steps=2) - - def test_training(self): - quantization_config = NVIDIAModelOptConfig(**self.get_dummy_init_kwargs()) - quantized_model = self.model_cls.from_pretrained( - self.model_id, - subfolder="transformer", - quantization_config=quantization_config, - torch_dtype=torch.bfloat16, - ).to(torch_device) - - for param in quantized_model.parameters(): - param.requires_grad = False - if param.ndim == 1: - param.data = param.data.to(torch.float32) - - for _, module in quantized_model.named_modules(): - if hasattr(module, "to_q"): - module.to_q = LoRALayer(module.to_q, rank=4) - if hasattr(module, "to_k"): - module.to_k = LoRALayer(module.to_k, rank=4) - if hasattr(module, "to_v"): - module.to_v = LoRALayer(module.to_v, rank=4) - - with torch.amp.autocast(str(torch_device), dtype=torch.bfloat16): - inputs = self.get_dummy_inputs() - output = quantized_model(**inputs)[0] - output.norm().backward() - - for module in quantized_model.modules(): - if isinstance(module, LoRALayer): - assert module.adapter[1].weight.grad is not None - - -class TestSanaTransformerFP8Weights(ModelOptBaseTesterMixin): - expected_memory_reduction = 0.6 - - def get_dummy_init_kwargs(self): - return {"quant_type": "FP8"} - - @require_modelopt_version_greater_or_equal("0.44.0") - def test_prequantized_serialization_with_device_map(self): - mto.enable_huggingface_checkpointing() - model = self.model_cls.from_pretrained( - self.model_id, - subfolder="transformer", - torch_dtype=self.torch_dtype, - quantization_config=NVIDIAModelOptConfig( - quant_type="FP8", modelopt_config=copy.deepcopy(mtq.FP8_DEFAULT_CFG) - ), - ) - model.to(torch_device) - - with tempfile.TemporaryDirectory() as tmp_dir: - model.save_pretrained(tmp_dir) - assert os.path.isfile(os.path.join(tmp_dir, "modelopt_state.pth")) - saved_model = self.model_cls.from_pretrained( - tmp_dir, - torch_dtype=self.torch_dtype, - device_map=torch_device, - ) - - named_parameters = list(saved_model.named_parameters()) - named_buffers = list(saved_model.named_buffers()) - assert any(name.endswith(("_amax", "_scale")) for name, _ in named_buffers), ( - "The restored model did not contain ModelOpt quantizer buffers." - ) - - for tensor_kind, named_tensors in (("parameter", named_parameters), ("buffer", named_buffers)): - for name, tensor in named_tensors: - assert not tensor.is_meta, f"{tensor_kind} {name} was not materialized from meta." - - -class TestSanaTransformerINT8Weights(ModelOptBaseTesterMixin): - expected_memory_reduction = 0.6 - _test_torch_compile = True - - def get_dummy_init_kwargs(self): - return {"quant_type": "INT8"} - - -@require_torch_cuda_compatibility(8.0) -class TestSanaTransformerINT4Weights(ModelOptBaseTesterMixin): - expected_memory_reduction = 0.55 - - def get_dummy_init_kwargs(self): - return { - "quant_type": "INT4", - "block_quantize": 128, - "channel_quantize": -1, - "disable_conv_quantization": True, - } - - -@require_torch_cuda_compatibility(8.0) -class TestSanaTransformerNF4Weights(ModelOptBaseTesterMixin): - expected_memory_reduction = 0.65 - - def get_dummy_init_kwargs(self): - return { - "quant_type": "NF4", - "block_quantize": 128, - "channel_quantize": -1, - "scale_block_quantize": 8, - "scale_channel_quantize": -1, - "modules_to_not_convert": ["conv"], - } - - -@require_torch_cuda_compatibility(8.0) -class TestSanaTransformerNVFP4Weights(ModelOptBaseTesterMixin): - expected_memory_reduction = 0.65 - - def get_dummy_init_kwargs(self): - return { - "quant_type": "NVFP4", - "block_quantize": 128, - "channel_quantize": -1, - "scale_block_quantize": 8, - "scale_channel_quantize": -1, - "modules_to_not_convert": ["conv"], - } diff --git a/tests/quantization/quanto/__init__.py b/tests/quantization/quanto/__init__.py deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/tests/quantization/quanto/test_quanto.py b/tests/quantization/quanto/test_quanto.py deleted file mode 100644 index 6b1ee57efae9..000000000000 --- a/tests/quantization/quanto/test_quanto.py +++ /dev/null @@ -1,176 +0,0 @@ -import gc - -import pytest - -from diffusers import FluxPipeline, FluxTransformer2DModel, QuantoConfig -from diffusers.utils import is_torch_available - -from ...testing_utils import ( - backend_empty_cache, - backend_reset_peak_memory_stats, - enable_full_determinism, - nightly, - require_accelerate, - require_accelerator, - require_torch_cuda_compatibility, - torch_device, -) - - -if is_torch_available(): - import torch - - from ..utils import get_memory_consumption_stat - -enable_full_determinism() - - -@nightly -@require_accelerator -@require_accelerate -class QuantoBaseTesterMixin: - model_id = None - pipeline_model_id = None - model_cls = None - torch_dtype = torch.bfloat16 - # the expected reduction in peak memory used compared to an unquantized model expressed as a percentage - expected_memory_reduction = 0.0 - - @pytest.fixture(autouse=True) - def _cleanup(self): - backend_reset_peak_memory_stats(torch_device) - backend_empty_cache(torch_device) - gc.collect() - yield - backend_reset_peak_memory_stats(torch_device) - backend_empty_cache(torch_device) - gc.collect() - - def get_dummy_init_kwargs(self): - return {"weights_dtype": "float8"} - - def get_dummy_model_init_kwargs(self): - return { - "pretrained_model_name_or_path": self.model_id, - "torch_dtype": self.torch_dtype, - "quantization_config": QuantoConfig(**self.get_dummy_init_kwargs()), - } - - def test_quanto_memory_usage(self): - inputs = self.get_dummy_inputs() - inputs = { - k: v.to(device=torch_device, dtype=torch.bfloat16) for k, v in inputs.items() if not isinstance(v, bool) - } - - unquantized_model = self.model_cls.from_pretrained(self.model_id, torch_dtype=self.torch_dtype) - unquantized_model.to(torch_device) - unquantized_model_memory = get_memory_consumption_stat(unquantized_model, inputs) - - quantized_model = self.model_cls.from_pretrained(**self.get_dummy_model_init_kwargs()) - quantized_model.to(torch_device) - quantized_model_memory = get_memory_consumption_stat(quantized_model, inputs) - - assert unquantized_model_memory / quantized_model_memory >= self.expected_memory_reduction - - def test_dtype_assignment(self): - model = self.model_cls.from_pretrained(**self.get_dummy_model_init_kwargs()) - - with pytest.raises(ValueError): - # Tries with a `dtype` - model.to(torch.float16) - - with pytest.raises(ValueError): - # Tries with a `device` and `dtype` - device_0 = f"{torch_device}:0" - model.to(device=device_0, dtype=torch.float16) - - with pytest.raises(ValueError): - # Tries with a cast - model.float() - - with pytest.raises(ValueError): - # Tries with a cast - model.half() - - # This should work - model.to(torch_device) - - def test_device_map_error(self): - with pytest.raises(ValueError): - _ = self.model_cls.from_pretrained( - **self.get_dummy_model_init_kwargs(), device_map={0: "8GB", "cpu": "16GB"} - ) - - -class FluxTransformerQuantoMixin(QuantoBaseTesterMixin): - model_id = "hf-internal-testing/tiny-flux-transformer" - model_cls = FluxTransformer2DModel - pipeline_cls = FluxPipeline - torch_dtype = torch.bfloat16 - - def get_dummy_inputs(self): - return { - "hidden_states": torch.randn((1, 4096, 64), generator=torch.Generator("cpu").manual_seed(0)).to( - torch_device, self.torch_dtype - ), - "encoder_hidden_states": torch.randn( - (1, 512, 4096), - generator=torch.Generator("cpu").manual_seed(0), - ).to(torch_device, self.torch_dtype), - "pooled_projections": torch.randn( - (1, 768), - generator=torch.Generator("cpu").manual_seed(0), - ).to(torch_device, self.torch_dtype), - "timestep": torch.tensor([1]).to(torch_device, self.torch_dtype), - "img_ids": torch.randn((4096, 3), generator=torch.Generator("cpu").manual_seed(0)).to( - torch_device, self.torch_dtype - ), - "txt_ids": torch.randn((512, 3), generator=torch.Generator("cpu").manual_seed(0)).to( - torch_device, self.torch_dtype - ), - "guidance": torch.tensor([3.5]).to(torch_device, self.torch_dtype), - } - - def test_model_cpu_offload(self): - init_kwargs = self.get_dummy_init_kwargs() - transformer = self.model_cls.from_pretrained( - "hf-internal-testing/tiny-flux-pipe", - quantization_config=QuantoConfig(**init_kwargs), - subfolder="transformer", - torch_dtype=torch.bfloat16, - ) - pipe = self.pipeline_cls.from_pretrained( - "hf-internal-testing/tiny-flux-pipe", transformer=transformer, torch_dtype=torch.bfloat16 - ) - pipe.enable_model_cpu_offload(device=torch_device) - _ = pipe("a cat holding a sign that says hello", num_inference_steps=2) - - -class TestFluxTransformerFloat8Weights(FluxTransformerQuantoMixin): - expected_memory_reduction = 0.6 - - def get_dummy_init_kwargs(self): - return {"weights_dtype": "float8"} - - -class TestFluxTransformerInt8Weights(FluxTransformerQuantoMixin): - expected_memory_reduction = 0.6 - - def get_dummy_init_kwargs(self): - return {"weights_dtype": "int8"} - - -@require_torch_cuda_compatibility(8.0) -class TestFluxTransformerInt4Weights(FluxTransformerQuantoMixin): - expected_memory_reduction = 0.55 - - def get_dummy_init_kwargs(self): - return {"weights_dtype": "int4"} - - -@require_torch_cuda_compatibility(8.0) -class TestFluxTransformerInt2Weights(FluxTransformerQuantoMixin): - expected_memory_reduction = 0.65 - - def get_dummy_init_kwargs(self): - return {"weights_dtype": "int2"} diff --git a/tests/quantization/test_pipeline_level_quantization.py b/tests/quantization/test_pipeline_level_quantization.py deleted file mode 100644 index 18679427f18e..000000000000 --- a/tests/quantization/test_pipeline_level_quantization.py +++ /dev/null @@ -1,312 +0,0 @@ -# coding=utf-8 -# Copyright 2025 The HuggingFace Team 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 clone 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 json -import tempfile - -import pytest -import torch -from parameterized import parameterized - -from diffusers import BitsAndBytesConfig, DiffusionPipeline, QuantoConfig -from diffusers.quantizers import PipelineQuantizationConfig -from diffusers.utils import logging - -from ..testing_utils import ( - CaptureLogger, - is_transformers_available, - require_accelerate, - require_bitsandbytes_version_greater, - require_quanto, - require_torch, - require_torch_accelerator, - slow, - torch_device, -) - - -if is_transformers_available(): - from transformers import BitsAndBytesConfig as TranBitsAndBytesConfig -else: - TranBitsAndBytesConfig = None - - -@require_bitsandbytes_version_greater("0.43.2") -@require_quanto -@require_accelerate -@require_torch -@require_torch_accelerator -@slow -class TestPipelineQuantization: - model_name = "hf-internal-testing/tiny-flux-pipe" - prompt = "a beautiful sunset amidst the mountains." - num_inference_steps = 10 - seed = 0 - - def test_quant_config_set_correctly_through_kwargs(self): - components_to_quantize = ["transformer", "text_encoder_2"] - quant_config = PipelineQuantizationConfig( - quant_backend="bitsandbytes_4bit", - quant_kwargs={ - "load_in_4bit": True, - "bnb_4bit_quant_type": "nf4", - "bnb_4bit_compute_dtype": torch.bfloat16, - }, - components_to_quantize=components_to_quantize, - ) - pipe = DiffusionPipeline.from_pretrained( - self.model_name, - quantization_config=quant_config, - torch_dtype=torch.bfloat16, - ).to(torch_device) - for name, component in pipe.components.items(): - if name in components_to_quantize: - assert getattr(component.config, "quantization_config", None) is not None - quantization_config = component.config.quantization_config - assert quantization_config.load_in_4bit - assert quantization_config.quant_method == "bitsandbytes" - - _ = pipe(self.prompt, num_inference_steps=self.num_inference_steps) - - def test_quant_config_set_correctly_through_granular(self): - quant_config = PipelineQuantizationConfig( - quant_mapping={ - "transformer": QuantoConfig(weights_dtype="int8"), - "text_encoder_2": TranBitsAndBytesConfig(load_in_4bit=True, compute_dtype=torch.bfloat16), - } - ) - components_to_quantize = list(quant_config.quant_mapping.keys()) - pipe = DiffusionPipeline.from_pretrained( - self.model_name, - quantization_config=quant_config, - torch_dtype=torch.bfloat16, - ).to(torch_device) - for name, component in pipe.components.items(): - if name in components_to_quantize: - assert getattr(component.config, "quantization_config", None) is not None - quantization_config = component.config.quantization_config - - if name == "text_encoder_2": - assert quantization_config.load_in_4bit - assert quantization_config.quant_method == "bitsandbytes" - else: - assert quantization_config.quant_method == "quanto" - - _ = pipe(self.prompt, num_inference_steps=self.num_inference_steps) - - def test_raises_error_for_invalid_config(self): - with pytest.raises(ValueError) as err_context: - _ = PipelineQuantizationConfig( - quant_mapping={ - "transformer": QuantoConfig(weights_dtype="int8"), - "text_encoder_2": TranBitsAndBytesConfig(load_in_4bit=True, compute_dtype=torch.bfloat16), - }, - quant_backend="bitsandbytes_4bit", - ) - - assert ( - str(err_context.value) == "Both `quant_backend` and `quant_mapping` cannot be specified at the same time." - ) - - def test_validation_for_kwargs(self): - components_to_quantize = ["transformer", "text_encoder_2"] - with pytest.raises(ValueError) as err_context: - _ = PipelineQuantizationConfig( - quant_backend="quanto", - quant_kwargs={"weights_dtype": "int8"}, - components_to_quantize=components_to_quantize, - ) - - assert "The signatures of the __init__ methods of the quantization config classes" in str(err_context.value) - - def test_raises_error_for_wrong_config_class(self): - quant_config = { - "transformer": QuantoConfig(weights_dtype="int8"), - "text_encoder_2": TranBitsAndBytesConfig(load_in_4bit=True, compute_dtype=torch.bfloat16), - } - with pytest.raises(ValueError) as err_context: - _ = DiffusionPipeline.from_pretrained( - self.model_name, - quantization_config=quant_config, - torch_dtype=torch.bfloat16, - ) - assert str(err_context.value) == "`quantization_config` must be an instance of `PipelineQuantizationConfig`." - - def test_validation_for_mapping(self): - with pytest.raises(ValueError) as err_context: - _ = PipelineQuantizationConfig( - quant_mapping={ - "transformer": DiffusionPipeline(), - "text_encoder_2": TranBitsAndBytesConfig(load_in_4bit=True, compute_dtype=torch.bfloat16), - } - ) - - assert "Provided config for module_name=transformer could not be found" in str(err_context.value) - - def test_saving_loading(self): - quant_config = PipelineQuantizationConfig( - quant_mapping={ - "transformer": QuantoConfig(weights_dtype="int8"), - "text_encoder_2": TranBitsAndBytesConfig(load_in_4bit=True, compute_dtype=torch.bfloat16), - } - ) - components_to_quantize = list(quant_config.quant_mapping.keys()) - pipe = DiffusionPipeline.from_pretrained( - self.model_name, - quantization_config=quant_config, - torch_dtype=torch.bfloat16, - ).to(torch_device) - - pipe_inputs = {"prompt": self.prompt, "num_inference_steps": self.num_inference_steps, "output_type": "latent"} - output_1 = pipe(**pipe_inputs, generator=torch.manual_seed(self.seed)).images - - with tempfile.TemporaryDirectory() as tmpdir: - pipe.save_pretrained(tmpdir) - loaded_pipe = DiffusionPipeline.from_pretrained(tmpdir, torch_dtype=torch.bfloat16).to(torch_device) - for name, component in loaded_pipe.components.items(): - if name in components_to_quantize: - assert getattr(component.config, "quantization_config", None) is not None - quantization_config = component.config.quantization_config - - if name == "text_encoder_2": - assert quantization_config.load_in_4bit - assert quantization_config.quant_method == "bitsandbytes" - else: - assert quantization_config.quant_method == "quanto" - - output_2 = loaded_pipe(**pipe_inputs, generator=torch.manual_seed(self.seed)).images - - assert torch.allclose(output_1, output_2) - - @parameterized.expand(["quant_kwargs", "quant_mapping"]) - def test_warn_invalid_component(self, method): - invalid_component = "foo" - if method == "quant_kwargs": - components_to_quantize = ["transformer", invalid_component] - quant_config = PipelineQuantizationConfig( - quant_backend="bitsandbytes_8bit", - quant_kwargs={"load_in_8bit": True}, - components_to_quantize=components_to_quantize, - ) - else: - quant_config = PipelineQuantizationConfig( - quant_mapping={ - "transformer": QuantoConfig("int8"), - invalid_component: TranBitsAndBytesConfig(load_in_8bit=True), - } - ) - - logger = logging.get_logger("diffusers.pipelines.pipeline_loading_utils") - logger.setLevel(logging.WARNING) - with CaptureLogger(logger) as cap_logger: - _ = DiffusionPipeline.from_pretrained( - self.model_name, - quantization_config=quant_config, - torch_dtype=torch.bfloat16, - ) - assert invalid_component in cap_logger.out - - @parameterized.expand(["quant_kwargs", "quant_mapping"]) - def test_no_quantization_for_all_invalid_components(self, method): - invalid_component = "foo" - if method == "quant_kwargs": - components_to_quantize = [invalid_component] - quant_config = PipelineQuantizationConfig( - quant_backend="bitsandbytes_8bit", - quant_kwargs={"load_in_8bit": True}, - components_to_quantize=components_to_quantize, - ) - else: - quant_config = PipelineQuantizationConfig( - quant_mapping={invalid_component: TranBitsAndBytesConfig(load_in_8bit=True)} - ) - - pipe = DiffusionPipeline.from_pretrained( - self.model_name, - quantization_config=quant_config, - torch_dtype=torch.bfloat16, - ) - for name, component in pipe.components.items(): - if isinstance(component, torch.nn.Module): - assert not hasattr(component.config, "quantization_config") - - @parameterized.expand(["quant_kwargs", "quant_mapping"]) - def test_quant_config_repr(self, method): - component_name = "transformer" - if method == "quant_kwargs": - components_to_quantize = [component_name] - quant_config = PipelineQuantizationConfig( - quant_backend="bitsandbytes_8bit", - quant_kwargs={"load_in_8bit": True}, - components_to_quantize=components_to_quantize, - ) - else: - quant_config = PipelineQuantizationConfig( - quant_mapping={component_name: BitsAndBytesConfig(load_in_8bit=True)} - ) - - pipe = DiffusionPipeline.from_pretrained( - self.model_name, - quantization_config=quant_config, - torch_dtype=torch.bfloat16, - ) - assert getattr(pipe, "quantization_config", None) is not None - retrieved_config = pipe.quantization_config - expected_config = """ -transformer BitsAndBytesConfig { - "_load_in_4bit": false, - "_load_in_8bit": true, - "bnb_4bit_compute_dtype": "float32", - "bnb_4bit_quant_storage": "uint8", - "bnb_4bit_quant_type": "fp4", - "bnb_4bit_use_double_quant": false, - "llm_int8_enable_fp32_cpu_offload": false, - "llm_int8_has_fp16_weight": false, - "llm_int8_skip_modules": null, - "llm_int8_threshold": 6.0, - "load_in_4bit": false, - "load_in_8bit": true, - "quant_method": "bitsandbytes" -} - -""" - expected_data = self._parse_config_string(expected_config) - actual_data = self._parse_config_string(str(retrieved_config)) - assert actual_data == expected_data - - def _parse_config_string(self, config_string: str) -> tuple[str, dict]: - first_brace = config_string.find("{") - if first_brace == -1: - raise ValueError("Could not find opening brace '{' in the string.") - - json_part = config_string[first_brace:] - data = json.loads(json_part) - - return data - - def test_single_component_to_quantize(self): - component_to_quantize = "transformer" - quant_config = PipelineQuantizationConfig( - quant_backend="bitsandbytes_8bit", - quant_kwargs={"load_in_8bit": True}, - components_to_quantize=component_to_quantize, - ) - pipe = DiffusionPipeline.from_pretrained( - self.model_name, - quantization_config=quant_config, - torch_dtype=torch.bfloat16, - ) - for name, component in pipe.components.items(): - if name == component_to_quantize: - assert hasattr(component.config, "quantization_config") diff --git a/tests/quantization/test_torch_compile_utils.py b/tests/quantization/test_torch_compile_utils.py deleted file mode 100644 index 9b800ffaa30b..000000000000 --- a/tests/quantization/test_torch_compile_utils.py +++ /dev/null @@ -1,66 +0,0 @@ -# coding=utf-8 -# Copyright 2025 The HuggingFace Team 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 clone 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 pytest -import torch - -from diffusers import DiffusionPipeline - -from ..testing_utils import backend_empty_cache, require_torch_accelerator, slow, torch_device - - -@require_torch_accelerator -@slow -class QuantCompileTests: - @property - def quantization_config(self): - raise NotImplementedError( - "This property should be implemented in the subclass to return the appropriate quantization config." - ) - - @pytest.fixture(autouse=True) - def _cleanup(self): - gc.collect() - backend_empty_cache(torch_device) - torch.compiler.reset() - yield - gc.collect() - backend_empty_cache(torch_device) - torch.compiler.reset() - - def _init_pipeline(self, quantization_config, torch_dtype): - pipe = DiffusionPipeline.from_pretrained( - "stabilityai/stable-diffusion-3-medium-diffusers", - quantization_config=quantization_config, - torch_dtype=torch_dtype, - ) - return pipe - - def _test_torch_compile_with_cpu_offload(self, torch_dtype=torch.bfloat16): - pipe = self._init_pipeline(self.quantization_config, torch_dtype) - pipe.enable_model_cpu_offload() - # regional compilation is better for offloading. - # see: https://pytorch.org/blog/torch-compile-and-diffusers-a-hands-on-guide-to-peak-performance/ - if getattr(pipe.transformer, "_repeated_blocks"): - pipe.transformer.compile_repeated_blocks(fullgraph=True) - else: - pipe.transformer.compile() - - # small resolutions to ensure speedy execution. - pipe("a dog", num_inference_steps=2, max_sequence_length=16, height=256, width=256) - - def test_torch_compile_with_cpu_offload(self): - self._test_torch_compile_with_cpu_offload() diff --git a/tests/quantization/torchao/README.md b/tests/quantization/torchao/README.md deleted file mode 100644 index 1b06be1b83e0..000000000000 --- a/tests/quantization/torchao/README.md +++ /dev/null @@ -1,50 +0,0 @@ -The tests here are adapted from [`transformers` tests](https://github.com/huggingface/transformers/blob/3a8eb74668e9c2cc563b2f5c62fac174797063e0/tests/quantization/torchao_integration/). - -The benchmarks were run on a single H100. Below is `nvidia-smi`: - -```bash -+---------------------------------------------------------------------------------------+ -| NVIDIA-SMI 535.104.12 Driver Version: 535.104.12 CUDA Version: 12.2 | -|-----------------------------------------+----------------------+----------------------+ -| GPU Name Persistence-M | Bus-Id Disp.A | Volatile Uncorr. ECC | -| Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. | -| | | MIG M. | -|=========================================+======================+======================| -| 0 NVIDIA H100 80GB HBM3 On | 00000000:53:00.0 Off | 0 | -| N/A 34C P0 69W / 700W | 2MiB / 81559MiB | 0% Default | -| | | Disabled | -+-----------------------------------------+----------------------+----------------------+ - -+---------------------------------------------------------------------------------------+ -| Processes: | -| GPU GI CI PID Type Process name GPU Memory | -| ID ID Usage | -|=======================================================================================| -| No running processes found | -+---------------------------------------------------------------------------------------+ -``` - -The benchmark results for Flux and CogVideoX can be found in [this](https://github.com/huggingface/diffusers/pull/10009) PR. - -The tests, and the expected slices, were obtained from the `aws-g6e-xlarge-plus` GPU test runners. To run the slow tests, use the following command or an equivalent: - -```bash -HF_XET_HIGH_PERFORMANCE=1 RUN_SLOW=1 pytest -s tests/quantization/torchao/test_torchao.py::SlowTorchAoTests -``` - -`diffusers-cli`: - -```bash -- 🤗 Diffusers version: 0.32.0.dev0 -- Platform: Linux-5.15.0-1049-aws-x86_64-with-glibc2.31 -- Running on Google Colab?: No -- Python version: 3.10.14 -- PyTorch version (GPU?): 2.6.0.dev20241112+cu121 (False) -- Huggingface_hub version: 0.26.2 -- Transformers version: 4.46.3 -- Accelerate version: 1.1.1 -- PEFT version: not installed -- Bitsandbytes version: not installed -- Safetensors version: 0.4.5 -- xFormers version: not installed -``` diff --git a/tests/quantization/torchao/test_torchao.py b/tests/quantization/torchao/test_torchao.py index ea2dad559b49..6666569187cb 100644 --- a/tests/quantization/torchao/test_torchao.py +++ b/tests/quantization/torchao/test_torchao.py @@ -13,79 +13,34 @@ # See the License for the specific language governing permissions and # limitations under the License. -import gc -import tempfile -from typing import List -import numpy as np import pytest -from transformers import AutoTokenizer, CLIPTextModel, CLIPTokenizer, T5EncoderModel -from diffusers import ( - AutoencoderKL, - FlowMatchEulerDiscreteScheduler, - FluxPipeline, - FluxTransformer2DModel, - TorchAoConfig, -) -from diffusers.quantizers import PipelineQuantizationConfig +from diffusers import TorchAoConfig from ...testing_utils import ( - Expectations, - backend_empty_cache, - backend_synchronize, enable_full_determinism, - is_torch_available, + is_quantization, + is_torchao, is_torchao_available, - nightly, - numpy_cosine_similarity_distance, require_torch, require_torch_accelerator, require_torchao_version_greater_or_equal, - slow, - torch_device, ) -from ..test_torch_compile_utils import QuantCompileTests - - -enable_full_determinism() - - -def _is_xpu_or_cuda_capability_atleast_8_9() -> bool: - if is_torch_available(): - import torch - - if torch.cuda.is_available(): - major, minor = torch.cuda.get_device_capability() - if major == 8: - return minor >= 9 - return major >= 9 - elif torch.xpu.is_available(): - return True - return False - -if is_torch_available(): - import torch - import torch.nn as nn - from ..utils import get_memory_consumption_stat +if is_torchao_available(): + from torchao.quantization import Int4WeightOnlyConfig, Int8WeightOnlyConfig -if is_torchao_available(): - from torchao.quantization import ( - Float8WeightOnlyConfig, - Int4Tensor, - Int4WeightOnlyConfig, - Int8DynamicActivationInt8WeightConfig, - Int8DynamicActivationIntxWeightConfig, - Int8Tensor, - Int8WeightOnlyConfig, - IntxWeightOnlyConfig, - ) - from torchao.utils import TorchAOBaseTensor, get_model_size_in_bytes +enable_full_determinism() +# Model-level TorchAO tests live in `tests/models/testing_utils/quantization.py` and +# pipeline-level ones in `tests/pipelines/testing_utils/quantization.py`. This module covers +# backend behavior that fits neither: config validation and custom device maps with cpu/disk offload. +@is_quantization +@is_torchao @require_torch @require_torch_accelerator @require_torchao_version_greater_or_equal("0.15.0") @@ -118,625 +73,3 @@ def test_repr(self): quantization_repr = repr(quantization_config) assert "TorchAoConfig" in quantization_repr assert "torchao" in quantization_repr - - -# Slices for these tests have been obtained on our aws-g6e-xlarge-plus runners -@require_torch -@require_torch_accelerator -@require_torchao_version_greater_or_equal("0.15.0") -class TestTorchAo: - @pytest.fixture(autouse=True) - def _setup_torchao(self): - yield - gc.collect() - backend_empty_cache(torch_device) - - def get_dummy_components( - self, quantization_config: TorchAoConfig, model_id: str = "hf-internal-testing/tiny-flux-pipe" - ): - transformer = FluxTransformer2DModel.from_pretrained( - model_id, - subfolder="transformer", - quantization_config=quantization_config, - torch_dtype=torch.bfloat16, - ) - text_encoder = CLIPTextModel.from_pretrained(model_id, subfolder="text_encoder", torch_dtype=torch.bfloat16) - text_encoder_2 = T5EncoderModel.from_pretrained( - model_id, subfolder="text_encoder_2", torch_dtype=torch.bfloat16 - ) - tokenizer = CLIPTokenizer.from_pretrained(model_id, subfolder="tokenizer") - tokenizer_2 = AutoTokenizer.from_pretrained(model_id, subfolder="tokenizer_2") - vae = AutoencoderKL.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.bfloat16) - scheduler = FlowMatchEulerDiscreteScheduler() - - return { - "scheduler": scheduler, - "text_encoder": text_encoder, - "text_encoder_2": text_encoder_2, - "tokenizer": tokenizer, - "tokenizer_2": tokenizer_2, - "transformer": transformer, - "vae": vae, - } - - def get_dummy_inputs(self, device: torch.device, seed: int = 0): - if str(device).startswith("mps"): - generator = torch.manual_seed(seed) - else: - generator = torch.Generator().manual_seed(seed) - - inputs = { - "prompt": "an astronaut riding a horse in space", - "height": 32, - "width": 32, - "num_inference_steps": 2, - "output_type": "np", - "generator": generator, - } - - return inputs - - def get_dummy_tensor_inputs(self, device=None, seed: int = 0): - batch_size = 1 - num_latent_channels = 4 - num_image_channels = 3 - height = width = 4 - sequence_length = 48 - embedding_dim = 32 - - torch.manual_seed(seed) - hidden_states = torch.randn((batch_size, height * width, num_latent_channels)).to(device, dtype=torch.bfloat16) - - torch.manual_seed(seed) - encoder_hidden_states = torch.randn((batch_size, sequence_length, embedding_dim)).to( - device, dtype=torch.bfloat16 - ) - - torch.manual_seed(seed) - pooled_prompt_embeds = torch.randn((batch_size, embedding_dim)).to(device, dtype=torch.bfloat16) - - torch.manual_seed(seed) - text_ids = torch.randn((sequence_length, num_image_channels)).to(device, dtype=torch.bfloat16) - - torch.manual_seed(seed) - image_ids = torch.randn((height * width, num_image_channels)).to(device, dtype=torch.bfloat16) - - timestep = torch.tensor([1.0]).to(device, dtype=torch.bfloat16).expand(batch_size) - - return { - "hidden_states": hidden_states, - "encoder_hidden_states": encoder_hidden_states, - "pooled_projections": pooled_prompt_embeds, - "txt_ids": text_ids, - "img_ids": image_ids, - "timestep": timestep, - } - - def _test_quant_type(self, quantization_config: TorchAoConfig, expected_slice: List[float], model_id: str): - components = self.get_dummy_components(quantization_config, model_id) - pipe = FluxPipeline(**components) - pipe.to(device=torch_device) - - inputs = self.get_dummy_inputs(torch_device) - output = pipe(**inputs)[0] - output_slice = output[-1, -1, -3:, -3:].flatten() - - assert np.allclose(output_slice, expected_slice, atol=1e-3, rtol=1e-3) - - def test_quantization(self): - for model_id in ["hf-internal-testing/tiny-flux-pipe", "hf-internal-testing/tiny-flux-sharded"]: - # fmt: off - QUANTIZATION_TYPES_TO_TEST = [ - (Int4WeightOnlyConfig(version=2), np.array([0.4648, 0.5234, 0.5547, 0.4219, 0.4414, 0.6445, 0.4336, 0.4531, 0.5625])), - (Int8DynamicActivationIntxWeightConfig(version=2), np.array([0.4688, 0.5195, 0.5547, 0.418, 0.4414, 0.6406, 0.4336, 0.4531, 0.5625])), - (Int8WeightOnlyConfig(version=2), np.array([0.4648, 0.5195, 0.5547, 0.4199, 0.4414, 0.6445, 0.4316, 0.4531, 0.5625])), - (Int8DynamicActivationInt8WeightConfig(version=2), np.array([0.4648, 0.5195, 0.5547, 0.4199, 0.4414, 0.6445, 0.4316, 0.4531, 0.5625])), - (IntxWeightOnlyConfig(dtype=torch.uint4, group_size=16, version=2), np.array([0.4609, 0.5234, 0.5508, 0.4199, 0.4336, 0.6406, 0.4316, 0.4531, 0.5625])), - (IntxWeightOnlyConfig(dtype=torch.uint7, group_size=16, version=2), np.array([0.4648, 0.5195, 0.5547, 0.4219, 0.4414, 0.6445, 0.4316, 0.4531, 0.5625])), - ] - - if _is_xpu_or_cuda_capability_atleast_8_9(): - QUANTIZATION_TYPES_TO_TEST.extend([ - (Float8WeightOnlyConfig(weight_dtype=torch.float8_e5m2), np.array([0.4590, 0.5273, 0.5547, 0.4219, 0.4375, 0.6406, 0.4316, 0.4512, 0.5625])), - (Float8WeightOnlyConfig(weight_dtype=torch.float8_e4m3fn), np.array([0.4648, 0.5234, 0.5547, 0.4219, 0.4414, 0.6406, 0.4316, 0.4531, 0.5625])), - ]) - # fmt: on - - for quant_config, expected_slice in QUANTIZATION_TYPES_TO_TEST: - quantization_config = TorchAoConfig(quant_type=quant_config, modules_to_not_convert=["x_embedder"]) - self._test_quant_type(quantization_config, expected_slice, model_id) - - def test_int4wo_quant_bfloat16_conversion(self): - """ - Tests whether the dtype of model will be modified to bfloat16 for int4 weight-only quantization. - """ - quantization_config = TorchAoConfig(Int4WeightOnlyConfig(group_size=64)) - quantized_model = FluxTransformer2DModel.from_pretrained( - "hf-internal-testing/tiny-flux-pipe", - subfolder="transformer", - quantization_config=quantization_config, - torch_dtype=torch.bfloat16, - device_map=f"{torch_device}:0", - ) - - weight = quantized_model.transformer_blocks[0].ff.net[2].weight - assert isinstance(weight, Int4Tensor) - - def test_device_map(self): - """ - Test if the quantized model int4 weight-only is working properly with "auto" and custom device maps. - The custom device map performs cpu/disk offloading as well. Also verifies that the device map is - correctly set (in the `hf_device_map` attribute of the model). - """ - custom_device_map_dict = { - "time_text_embed": torch_device, - "context_embedder": torch_device, - "x_embedder": torch_device, - "transformer_blocks.0": "cpu", - "single_transformer_blocks.0": "disk", - "norm_out": torch_device, - "proj_out": "cpu", - } - device_maps = ["auto", custom_device_map_dict] - - inputs = self.get_dummy_tensor_inputs(torch_device) - # requires with different expected slices since models are different due to offload (we don't quantize modules offloaded to cpu/disk) - expected_slice_auto = np.array( - [ - 0.34179688, - -0.03613281, - 0.01428223, - -0.22949219, - -0.49609375, - 0.4375, - -0.1640625, - -0.66015625, - 0.43164062, - ] - ) - expected_slice_offload = np.array( - [0.34375, -0.03515625, 0.0123291, -0.22753906, -0.49414062, 0.4375, -0.16308594, -0.66015625, 0.43554688] - ) - for device_map in device_maps: - if device_map == "auto": - expected_slice = expected_slice_auto - else: - expected_slice = expected_slice_offload - with tempfile.TemporaryDirectory() as offload_folder: - quantization_config = TorchAoConfig(Int4WeightOnlyConfig(group_size=64)) - quantized_model = FluxTransformer2DModel.from_pretrained( - "hf-internal-testing/tiny-flux-pipe", - subfolder="transformer", - quantization_config=quantization_config, - device_map=device_map, - torch_dtype=torch.bfloat16, - offload_folder=offload_folder, - ) - - weight = quantized_model.transformer_blocks[0].ff.net[2].weight - - # Note that when performing cpu/disk offload, the offloaded weights are not quantized, only the weights on the gpu. - # This is not the case when the model are already quantized - if "transformer_blocks.0" in device_map: - assert isinstance(weight, nn.Parameter) - else: - assert isinstance(weight, Int4Tensor) - - output = quantized_model(**inputs)[0] - output_slice = output.flatten()[-9:].detach().float().cpu().numpy() - assert numpy_cosine_similarity_distance(output_slice, expected_slice) < 2e-3 - - with tempfile.TemporaryDirectory() as offload_folder: - quantization_config = TorchAoConfig(Int4WeightOnlyConfig(group_size=64)) - quantized_model = FluxTransformer2DModel.from_pretrained( - "hf-internal-testing/tiny-flux-sharded", - subfolder="transformer", - quantization_config=quantization_config, - device_map=device_map, - torch_dtype=torch.bfloat16, - offload_folder=offload_folder, - ) - - weight = quantized_model.transformer_blocks[0].ff.net[2].weight - if "transformer_blocks.0" in device_map: - assert isinstance(weight, nn.Parameter) - else: - assert isinstance(weight, Int4Tensor) - - output = quantized_model(**inputs)[0] - output_slice = output.flatten()[-9:].detach().float().cpu().numpy() - assert numpy_cosine_similarity_distance(output_slice, expected_slice) < 2e-3 - - def test_memory_footprint(self): - r""" - A simple test to check if the model conversion has been done correctly by checking on the - memory footprint of the converted model and the class type of the linear layers of the converted models - """ - for model_id in ["hf-internal-testing/tiny-flux-pipe", "hf-internal-testing/tiny-flux-sharded"]: - transformer_int4wo = self.get_dummy_components(TorchAoConfig(Int4WeightOnlyConfig()), model_id=model_id)[ - "transformer" - ] - transformer_int4wo_gs32 = self.get_dummy_components( - TorchAoConfig(Int4WeightOnlyConfig(group_size=32)), model_id=model_id - )["transformer"] - transformer_int8wo = self.get_dummy_components(TorchAoConfig(Int8WeightOnlyConfig()), model_id=model_id)[ - "transformer" - ] - transformer_bf16 = self.get_dummy_components(None, model_id=model_id)["transformer"] - - # Will not quantized all the layers by default due to the model weights shapes not being divisible by group_size=64 - for block in transformer_int4wo.transformer_blocks: - assert isinstance(block.ff.net[2].weight, Int4Tensor) - assert isinstance(block.ff_context.net[2].weight, Int4Tensor) - - # Will quantize all the linear layers except x_embedder - for name, module in transformer_int4wo_gs32.named_modules(): - if isinstance(module, nn.Linear) and name not in ["x_embedder"]: - assert isinstance(module.weight, Int4Tensor) - - # Will quantize all the linear layers - for module in transformer_int8wo.modules(): - if isinstance(module, nn.Linear): - assert isinstance(module.weight, Int8Tensor) - - total_int4wo = get_model_size_in_bytes(transformer_int4wo) - total_int4wo_gs32 = get_model_size_in_bytes(transformer_int4wo_gs32) - total_int8wo = get_model_size_in_bytes(transformer_int8wo) - total_bf16 = get_model_size_in_bytes(transformer_bf16) - - # TODO: refactor to align with other quantization tests - # Latter has smaller group size, so more groups -> more scales and zero points - assert total_int4wo < total_int4wo_gs32 - # int8 quantizes more layers compare to int4 with default group size - assert total_int8wo < total_int4wo - # int4wo does not quantize too many layers because of default group size, but for the layers it does - # there is additional overhead of scales and zero points - assert total_bf16 < total_int4wo - - def test_model_memory_usage(self): - model_id = "hf-internal-testing/tiny-flux-pipe" - expected_memory_saving_ratios = Expectations( - { - # XPU: For this tiny model, per-tensor overheads (alignment, fragmentation, metadata) become visible. - # While XPU doesn't have the large fixed cuBLAS workspace of A100, these small overheads prevent reaching the ideal 2.0 ratio. - # Observed ~1.27x (158k vs 124k) for model size. - # The runtime memory overhead is ~88k for both bf16 and int8wo. Adding this to model size: (158k+88k)/(124k+88k) ≈ 1.15. - ("xpu", None): 1.15, - # On Ampere, the cuBLAS kernels used for matrix multiplication often allocate a fixed-size workspace. - # Since the tiny-flux model weights are likely smaller than or comparable to this workspace, the total memory is dominated by the workspace. - ("cuda", 8): 1.02, - # On Hopper, TorchAO utilizes newer, highly optimized kernels (via Triton or CUTLASS 3.x) that are designed to be workspace-free or use negligible extra memory. - # Additionally, Triton kernels often handle unaligned memory better, avoiding the padding overhead seen on other backends for tiny tensors. - # This allows it to achieve the near-ideal 2.0x compression ratio. - ("cuda", 9): 2.0, - } - ) - expected_memory_saving_ratio = expected_memory_saving_ratios.get_expectation() - inputs = self.get_dummy_tensor_inputs(device=torch_device) - - transformer_bf16 = self.get_dummy_components(None, model_id=model_id)["transformer"] - transformer_bf16.to(torch_device) - unquantized_model_memory = get_memory_consumption_stat(transformer_bf16, inputs) - del transformer_bf16 - - transformer_int8wo = self.get_dummy_components(TorchAoConfig(Int8WeightOnlyConfig()), model_id=model_id)[ - "transformer" - ] - transformer_int8wo.to(torch_device) - quantized_model_memory = get_memory_consumption_stat(transformer_int8wo, inputs) - assert unquantized_model_memory / quantized_model_memory >= expected_memory_saving_ratio - - def test_wrong_config(self): - with pytest.raises(TypeError): - self.get_dummy_components(TorchAoConfig("int42")) - - def test_sequential_cpu_offload(self): - r""" - A test that checks if inference runs as expected when sequential cpu offloading is enabled. - """ - quantization_config = TorchAoConfig(Int8WeightOnlyConfig()) - components = self.get_dummy_components(quantization_config) - pipe = FluxPipeline(**components) - pipe.enable_sequential_cpu_offload() - - inputs = self.get_dummy_inputs(torch_device) - _ = pipe(**inputs) - - @require_torchao_version_greater_or_equal("0.15.0") - def test_aobase_config(self): - quantization_config = TorchAoConfig(Int8WeightOnlyConfig()) - components = self.get_dummy_components(quantization_config) - pipe = FluxPipeline(**components).to(torch_device) - - inputs = self.get_dummy_inputs(torch_device) - _ = pipe(**inputs) - - -# Slices for these tests have been obtained on our aws-g6e-xlarge-plus runners -@require_torch -@require_torch_accelerator -@require_torchao_version_greater_or_equal("0.15.0") -class TestTorchAoSerialization: - model_name = "hf-internal-testing/tiny-flux-pipe" - - @pytest.fixture(autouse=True) - def _setup_torchao_serialization(self): - yield - gc.collect() - backend_empty_cache(torch_device) - - def get_dummy_model(self, quant_type, device=None): - quantization_config = TorchAoConfig(quant_type) - quantized_model = FluxTransformer2DModel.from_pretrained( - self.model_name, - subfolder="transformer", - quantization_config=quantization_config, - torch_dtype=torch.bfloat16, - ) - return quantized_model.to(device) - - def get_dummy_tensor_inputs(self, device=None, seed: int = 0): - batch_size = 1 - num_latent_channels = 4 - num_image_channels = 3 - height = width = 4 - sequence_length = 48 - embedding_dim = 32 - - torch.manual_seed(seed) - hidden_states = torch.randn((batch_size, height * width, num_latent_channels)).to(device, dtype=torch.bfloat16) - encoder_hidden_states = torch.randn((batch_size, sequence_length, embedding_dim)).to( - device, dtype=torch.bfloat16 - ) - pooled_prompt_embeds = torch.randn((batch_size, embedding_dim)).to(device, dtype=torch.bfloat16) - text_ids = torch.randn((sequence_length, num_image_channels)).to(device, dtype=torch.bfloat16) - image_ids = torch.randn((height * width, num_image_channels)).to(device, dtype=torch.bfloat16) - timestep = torch.tensor([1.0]).to(device, dtype=torch.bfloat16).expand(batch_size) - - return { - "hidden_states": hidden_states, - "encoder_hidden_states": encoder_hidden_states, - "pooled_projections": pooled_prompt_embeds, - "txt_ids": text_ids, - "img_ids": image_ids, - "timestep": timestep, - } - - def _test_original_model_expected_slice(self, quant_type, expected_slice): - quantized_model = self.get_dummy_model(quant_type, torch_device) - inputs = self.get_dummy_tensor_inputs(torch_device) - output = quantized_model(**inputs)[0] - output_slice = output.flatten()[-9:].detach().float().cpu().numpy() - weight = quantized_model.transformer_blocks[0].ff.net[2].weight - assert isinstance(weight, TorchAOBaseTensor) - assert numpy_cosine_similarity_distance(output_slice, expected_slice) < 1e-3 - - def _check_serialization_expected_slice(self, quant_type, expected_slice, device): - quantized_model = self.get_dummy_model(quant_type, device) - - with tempfile.TemporaryDirectory() as tmp_dir: - quantized_model.save_pretrained(tmp_dir, safe_serialization=False) - loaded_quantized_model = FluxTransformer2DModel.from_pretrained( - tmp_dir, torch_dtype=torch.bfloat16, use_safetensors=False - ).to(device=torch_device) - - inputs = self.get_dummy_tensor_inputs(torch_device) - output = loaded_quantized_model(**inputs)[0] - - output_slice = output.flatten()[-9:].detach().float().cpu().numpy() - assert isinstance(loaded_quantized_model.proj_out.weight, TorchAOBaseTensor) - assert numpy_cosine_similarity_distance(output_slice, expected_slice) < 1e-3 - - def test_int_a8w8_accelerator(self): - quant_type = Int8DynamicActivationInt8WeightConfig() - expected_slice = np.array([0.3633, -0.1357, -0.0188, -0.249, -0.4688, 0.5078, -0.1289, -0.6914, 0.4551]) - device = torch_device - self._test_original_model_expected_slice(quant_type, expected_slice) - self._check_serialization_expected_slice(quant_type, expected_slice, device) - - def test_int_a16w8_accelerator(self): - quant_type = Int8WeightOnlyConfig() - expected_slice = np.array([0.3613, -0.127, -0.0223, -0.2539, -0.459, 0.4961, -0.1357, -0.6992, 0.4551]) - device = torch_device - self._test_original_model_expected_slice(quant_type, expected_slice) - self._check_serialization_expected_slice(quant_type, expected_slice, device) - - def test_int_a8w8_cpu(self): - quant_type = Int8DynamicActivationInt8WeightConfig() - expected_slice = np.array([0.3633, -0.1357, -0.0188, -0.249, -0.4688, 0.5078, -0.1289, -0.6914, 0.4551]) - device = "cpu" - self._test_original_model_expected_slice(quant_type, expected_slice) - self._check_serialization_expected_slice(quant_type, expected_slice, device) - - def test_int_a16w8_cpu(self): - quant_type = Int8WeightOnlyConfig() - expected_slice = np.array([0.3613, -0.127, -0.0223, -0.2539, -0.459, 0.4961, -0.1357, -0.6992, 0.4551]) - device = "cpu" - self._test_original_model_expected_slice(quant_type, expected_slice) - self._check_serialization_expected_slice(quant_type, expected_slice, device) - - def test_aobase_config(self): - quant_type = Int8WeightOnlyConfig() - expected_slice = np.array([0.3613, -0.127, -0.0223, -0.2539, -0.459, 0.4961, -0.1357, -0.6992, 0.4551]) - device = torch_device - self._test_original_model_expected_slice(quant_type, expected_slice) - self._check_serialization_expected_slice(quant_type, expected_slice, device) - - -@require_torchao_version_greater_or_equal("0.15.0") -class TestTorchAoCompile(QuantCompileTests): - @property - def quantization_config(self): - return PipelineQuantizationConfig( - quant_mapping={"transformer": TorchAoConfig(Int8WeightOnlyConfig())}, - ) - - def test_torch_compile_with_cpu_offload(self): - pipe = self._init_pipeline(self.quantization_config, torch.bfloat16) - pipe.enable_model_cpu_offload() - # No compilation because it fails with: - # RuntimeError: _apply(): Couldn't swap Linear.weight - - # small resolutions to ensure speedy execution. - pipe("a dog", num_inference_steps=2, max_sequence_length=16, height=256, width=256) - - -# Slices for these tests have been obtained on our aws-g6e-xlarge-plus runners -@require_torch -@require_torch_accelerator -@require_torchao_version_greater_or_equal("0.15.0") -@slow -@nightly -class TestSlowTorchAo: - @pytest.fixture(autouse=True) - def _setup_slow_torchao(self): - yield - gc.collect() - backend_empty_cache(torch_device) - - def get_dummy_components(self, quantization_config: TorchAoConfig): - # This is just for convenience, so that we can modify it at one place for custom environments and locally testing - cache_dir = None - model_id = "black-forest-labs/FLUX.1-dev" - transformer = FluxTransformer2DModel.from_pretrained( - model_id, - subfolder="transformer", - quantization_config=quantization_config, - torch_dtype=torch.bfloat16, - cache_dir=cache_dir, - ) - text_encoder = CLIPTextModel.from_pretrained( - model_id, subfolder="text_encoder", torch_dtype=torch.bfloat16, cache_dir=cache_dir - ) - text_encoder_2 = T5EncoderModel.from_pretrained( - model_id, subfolder="text_encoder_2", torch_dtype=torch.bfloat16, cache_dir=cache_dir - ) - tokenizer = CLIPTokenizer.from_pretrained(model_id, subfolder="tokenizer", cache_dir=cache_dir) - tokenizer_2 = AutoTokenizer.from_pretrained(model_id, subfolder="tokenizer_2", cache_dir=cache_dir) - vae = AutoencoderKL.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.bfloat16, cache_dir=cache_dir) - scheduler = FlowMatchEulerDiscreteScheduler() - - return { - "scheduler": scheduler, - "text_encoder": text_encoder, - "text_encoder_2": text_encoder_2, - "tokenizer": tokenizer, - "tokenizer_2": tokenizer_2, - "transformer": transformer, - "vae": vae, - } - - def get_dummy_inputs(self, device: torch.device, seed: int = 0): - if str(device).startswith("mps"): - generator = torch.manual_seed(seed) - else: - generator = torch.Generator().manual_seed(seed) - - inputs = { - "prompt": "an astronaut riding a horse in space", - "height": 512, - "width": 512, - "num_inference_steps": 20, - "output_type": "np", - "generator": generator, - } - - return inputs - - def _test_quant_type(self, quantization_config, expected_slice): - components = self.get_dummy_components(quantization_config) - pipe = FluxPipeline(**components) - pipe.enable_model_cpu_offload() - - weight = pipe.transformer.transformer_blocks[0].ff.net[2].weight - assert isinstance(weight, TorchAOBaseTensor) - - inputs = self.get_dummy_inputs(torch_device) - output = pipe(**inputs)[0].flatten() - output_slice = np.concatenate((output[:16], output[-16:])) - assert np.allclose(output_slice, expected_slice, atol=1e-3, rtol=1e-3) - - def test_quantization(self): - # fmt: off - QUANTIZATION_TYPES_TO_TEST = [ - (Int8WeightOnlyConfig(), np.array([0.0505, 0.0742, 0.1367, 0.0429, 0.0585, 0.1386, 0.0585, 0.0703, 0.1367, 0.0566, 0.0703, 0.1464, 0.0546, 0.0703, 0.1425, 0.0546, 0.3535, 0.7578, 0.5000, 0.4062, 0.7656, 0.5117, 0.4121, 0.7656, 0.5117, 0.3984, 0.7578, 0.5234, 0.4023, 0.7382, 0.5390, 0.4570])), - (Int8DynamicActivationInt8WeightConfig(), np.array([0.0546, 0.0761, 0.1386, 0.0488, 0.0644, 0.1425, 0.0605, 0.0742, 0.1406, 0.0625, 0.0722, 0.1523, 0.0625, 0.0742, 0.1503, 0.0605, 0.3886, 0.7968, 0.5507, 0.4492, 0.7890, 0.5351, 0.4316, 0.8007, 0.5390, 0.4179, 0.8281, 0.5820, 0.4531, 0.7812, 0.5703, 0.4921])), - ] - - if _is_xpu_or_cuda_capability_atleast_8_9(): - QUANTIZATION_TYPES_TO_TEST.extend([ - (Float8WeightOnlyConfig(weight_dtype=torch.float8_e4m3fn), np.array([0.0546, 0.0722, 0.1328, 0.0468, 0.0585, 0.1367, 0.0605, 0.0703, 0.1328, 0.0625, 0.0703, 0.1445, 0.0585, 0.0703, 0.1406, 0.0605, 0.3496, 0.7109, 0.4843, 0.4042, 0.7226, 0.5000, 0.4160, 0.7031, 0.4824, 0.3886, 0.6757, 0.4667, 0.3710, 0.6679, 0.4902, 0.4238])), - ]) - # fmt: on - - for quant_config, expected_slice in QUANTIZATION_TYPES_TO_TEST: - quantization_config = TorchAoConfig(quant_type=quant_config, modules_to_not_convert=["x_embedder"]) - self._test_quant_type(quantization_config, expected_slice) - gc.collect() - backend_empty_cache(torch_device) - backend_synchronize(torch_device) - - def test_serialization_int8wo(self): - quantization_config = TorchAoConfig(Int8WeightOnlyConfig()) - components = self.get_dummy_components(quantization_config) - pipe = FluxPipeline(**components) - pipe.enable_model_cpu_offload() - - weight = pipe.transformer.x_embedder.weight - assert isinstance(weight, Int8Tensor) - - inputs = self.get_dummy_inputs(torch_device) - output = pipe(**inputs)[0].flatten()[:128] - - with tempfile.TemporaryDirectory() as tmp_dir: - pipe.transformer.save_pretrained(tmp_dir, safe_serialization=False) - pipe.remove_all_hooks() - del pipe.transformer - gc.collect() - backend_empty_cache(torch_device) - backend_synchronize(torch_device) - transformer = FluxTransformer2DModel.from_pretrained( - tmp_dir, torch_dtype=torch.bfloat16, use_safetensors=False - ) - pipe.transformer = transformer - pipe.enable_model_cpu_offload() - - weight = transformer.x_embedder.weight - assert isinstance(weight, Int8Tensor) - - loaded_output = pipe(**inputs)[0].flatten()[:128] - # Seems to require higher tolerance depending on which machine it is being run. - # A difference of 0.06 in normalized pixel space (-1 to 1), corresponds to a difference of - # 0.06 / 2 * 255 = 7.65 in pixel space (0 to 255). On our CI runners, the difference is about 0.04, - # on DGX it is 0.06, and on audace it is 0.037. So, we are using a tolerance of 0.06 here. - assert np.allclose(output, loaded_output, atol=0.06) - - def test_memory_footprint_int4wo(self): - # The original checkpoints are in bf16 and about 24 GB - expected_memory_in_gb = 6.0 - quantization_config = TorchAoConfig(Int4WeightOnlyConfig()) - cache_dir = None - transformer = FluxTransformer2DModel.from_pretrained( - "black-forest-labs/FLUX.1-dev", - subfolder="transformer", - quantization_config=quantization_config, - torch_dtype=torch.bfloat16, - cache_dir=cache_dir, - ) - int4wo_memory_in_gb = get_model_size_in_bytes(transformer) / 1024**3 - assert int4wo_memory_in_gb < expected_memory_in_gb - - def test_memory_footprint_int8wo(self): - # The original checkpoints are in bf16 and about 24 GB - expected_memory_in_gb = 12.0 - quantization_config = TorchAoConfig(Int8WeightOnlyConfig()) - cache_dir = None - transformer = FluxTransformer2DModel.from_pretrained( - "black-forest-labs/FLUX.1-dev", - subfolder="transformer", - quantization_config=quantization_config, - torch_dtype=torch.bfloat16, - cache_dir=cache_dir, - ) - int8wo_memory_in_gb = get_model_size_in_bytes(transformer) / 1024**3 - assert int8wo_memory_in_gb < expected_memory_in_gb diff --git a/tests/quantization/utils.py b/tests/quantization/utils.py deleted file mode 100644 index a74ece5a3a3a..000000000000 --- a/tests/quantization/utils.py +++ /dev/null @@ -1,45 +0,0 @@ -from diffusers.utils import is_torch_available - -from ..testing_utils import ( - backend_empty_cache, - backend_max_memory_allocated, - backend_reset_peak_memory_stats, - torch_device, -) - - -if is_torch_available(): - import torch - import torch.nn as nn - - class LoRALayer(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: nn.Module, rank: int): - super().__init__() - self.module = module - self.adapter = nn.Sequential( - nn.Linear(module.in_features, rank, bias=False), - nn.Linear(rank, module.out_features, bias=False), - ) - small_std = (2.0 / (5 * min(module.in_features, module.out_features))) ** 0.5 - nn.init.normal_(self.adapter[0].weight, std=small_std) - 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) - - @torch.no_grad() - @torch.inference_mode() - def get_memory_consumption_stat(model, inputs): - backend_reset_peak_memory_stats(torch_device) - backend_empty_cache(torch_device) - - model(**inputs) - max_mem_allocated = backend_max_memory_allocated(torch_device) - return max_mem_allocated From d5b6437f2bdf3226991e393bb167adf21501cdf0 Mon Sep 17 00:00:00 2001 From: "sayak@huggingface.co" Date: Mon, 10 Aug 2026 09:32:16 +0000 Subject: [PATCH 6/6] [ci] select nightly quantization tests by marker across test tiers Each backend job now runs `pytest -m ` over tests/models, tests/quantization, and tests/pipelines/testing_utils/quantization.py, giving the model-level mixin tests a nightly home with the backend dependencies installed. The torchao job additionally installs mslk. Co-Authored-By: Claude Fable 5 --- .github/workflows/nightly_tests.yml | 20 ++++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/.github/workflows/nightly_tests.yml b/.github/workflows/nightly_tests.yml index 678d106a5fc7..b285edb3e01e 100644 --- a/.github/workflows/nightly_tests.yml +++ b/.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: |