# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import json import re from pathlib import Path from typing import NamedTuple import pytest from _test_utils.examples.models import FLUX_SCHNELL_PATH from _test_utils.examples.run_command import run_example_command from _test_utils.torch.misc import minimum_sm from safetensors import safe_open from test_diffusers import CALIB_PROMPTS_FILE, assert_hf_ckpt_exported # Run the example scripts in this process: ~20s of startup saved per step (see conftest). pytestmark = pytest.mark.usefixtures("in_process_examples") class DiffuserHfExportModel(NamedTuple): name: str path: str dtype: str format_type: str quant_algo: str collect_method: str model_dtype: str = "Half" def quantize_and_export_hf(self, tmp_path: Path) -> Path: hf_ckpt_dir = tmp_path / f"{self.name}_{self.format_type}_hf_ckpt" cmd_args = [ "python", "quantize.py", "--model", self.name, "--override-model-path", self.path, "--calib-size", "4", "--batch-size", "2", "--n-steps", "2", "--percentile", "1.0", "--alpha", "0.8", "--format", self.format_type, "--quant-algo", self.quant_algo, "--collect-method", self.collect_method, "--model-dtype", self.model_dtype, "--trt-high-precision-dtype", self.dtype, "--prompts-file", str(CALIB_PROMPTS_FILE), "--hf-ckpt-dir", str(hf_ckpt_dir), ] run_example_command(cmd_args, "diffusers/quantization") return hf_ckpt_dir # The SDXL int8/fp8 and Wan 2.2 transformer int8/fp8 exports run in test_diffusers.py, from the # quantize run those tests already do (``hf_export=True``). @pytest.mark.parametrize( "model", [ DiffuserHfExportModel( name="flux-schnell", path=FLUX_SCHNELL_PATH, dtype="BFloat16", format_type="int8", quant_algo="smoothquant", collect_method="min-mean", model_dtype="BFloat16", ), pytest.param( DiffuserHfExportModel( name="flux-schnell", path=FLUX_SCHNELL_PATH, dtype="BFloat16", format_type="fp4", quant_algo="max", collect_method="default", model_dtype="BFloat16", ), marks=minimum_sm(89), ), ], ids=[ "flux_schnell_int8_smoothquant_min_mean", "flux_schnell_fp4_max_default", ], ) def test_diffusers_hf_ckpt_export(model: DiffuserHfExportModel, tmp_path: Path) -> None: assert_hf_ckpt_exported(model.quantize_and_export_hf(tmp_path)) class QwenHfExportModel(NamedTuple): format_type: str quant_algo: str is_svdquant: bool def quantize_and_export_hf(self, tiny_qwen_image_path: str, tmp_path: Path) -> Path: hf_ckpt_dir = tmp_path / f"qwen_{self.format_type}_{self.quant_algo}_hf_ckpt" cmd_args = [ "python", "quantize.py", "--model", "qwen-image", "--override-model-path", str(tiny_qwen_image_path), "--format", self.format_type, "--quant-algo", self.quant_algo, "--collect-method", "default", "--model-dtype", "BFloat16", "--trt-high-precision-dtype", "BFloat16", "--calib-size", "2", "--batch-size", "1", "--n-steps", "2", "--prompts-file", str(CALIB_PROMPTS_FILE), "--hf-ckpt-dir", str(hf_ckpt_dir), ] if self.is_svdquant: cmd_args.extend(["--lowrank", "8"]) run_example_command(cmd_args, "diffusers/quantization") return hf_ckpt_dir def _module_prefixes(keys: set[str], suffix: str) -> set[str]: """Module paths (key minus suffix) for every key ending in ``suffix``.""" return {k[: -len(suffix)] for k in keys if k.endswith(suffix)} def _block_indices(prefixes: set[str]) -> set[int]: """transformer_blocks indices referenced by a set of module prefixes.""" indices = set() for prefix in prefixes: match = re.search(r"transformer_blocks\.(\d+)\.", prefix) if match: indices.add(int(match.group(1))) return indices # Tiny Qwen fixture has 6 transformer blocks; the recipe excludes the first 2 and # last 2, so only blocks 2 and 3 are quantized. _QWEN_QUANTIZED_BLOCKS = {2, 3} _QWEN_LORA_RANK = 8 # Per quantized block: image-stream linears keep full SVDQuant (low-rank branch + # pre_quant_scale), while the text-stream and modulation linears match the # "svdquant_skip_layers" patterns for qwen-image in # examples/diffusers/quantization/models_utils.py and are exported as plain NVFP4. _QWEN_SVDQUANT_PROMOTED_SUFFIXES = ( ".attn.to_q", ".attn.to_k", ".attn.to_v", ".attn.to_out.0", ".img_mlp.net.0.proj", ".img_mlp.net.2", ) _QWEN_SVDQUANT_SKIPPED_SUFFIXES = ( ".attn.add_q_proj", ".attn.add_k_proj", ".attn.add_v_proj", ".attn.to_add_out", ".txt_mlp.net.0.proj", ".txt_mlp.net.2", ".img_mod.1", ".txt_mod.1", ) @pytest.mark.parametrize( "qwen_model", [ pytest.param(QwenHfExportModel("fp8", "max", False), marks=minimum_sm(89)), pytest.param(QwenHfExportModel("fp4", "max", False), marks=minimum_sm(89)), pytest.param(QwenHfExportModel("fp4", "svdquant", True), marks=minimum_sm(89)), ], ids=["qwen_fp8_max", "qwen_nvfp4_max", "qwen_nvfp4_svdquant"], ) def test_qwen_image_hf_ckpt_export( qwen_model: QwenHfExportModel, tiny_qwen_image_path: str, tmp_path: Path ) -> None: hf_ckpt_dir = qwen_model.quantize_and_export_hf(tiny_qwen_image_path, tmp_path) assert hf_ckpt_dir.exists(), f"HF checkpoint directory was not created: {hf_ckpt_dir}" # The transformer is the quantized component. transformer_dir = hf_ckpt_dir / "transformer" config_path = transformer_dir / "config.json" assert config_path.exists(), f"no transformer/config.json in {hf_ckpt_dir}" quant_config = json.loads(config_path.read_text()).get("quantization_config") assert quant_config is not None, "missing quantization_config" assert quant_config.get("quant_method") == "modelopt" keys: set[str] = set() lora_tensors: dict[str, object] = {} safetensors_files = sorted(transformer_dir.rglob("*.safetensors")) assert safetensors_files, f"no safetensors in {transformer_dir}" for path in safetensors_files: with safe_open(str(path), framework="pt") as handle: for key in handle.keys(): # noqa: SIM118 - safe_open is not iterable keys.add(key) if key.endswith((".svdquant_lora_a", ".svdquant_lora_b")): lora_tensors[key] = handle.get_tensor(key) # No live quantizer state should leak into the exported checkpoint. assert not any("weight_quantizer" in k for k in keys), "quantizer keys leaked into export" assert not any("input_quantizer._amax" in k for k in keys) # Recipe: only the middle transformer blocks are quantized — first-2/last-2 of # transformer_blocks are excluded, and nothing outside transformer_blocks. weight_scale_prefixes = _module_prefixes(keys, ".weight_scale") assert weight_scale_prefixes, "no quantized linears found in export" assert all("transformer_blocks." in p for p in weight_scale_prefixes), ( f"a non-transformer_blocks module was quantized: {weight_scale_prefixes}" ) assert _block_indices(weight_scale_prefixes) == _QWEN_QUANTIZED_BLOCKS, ( f"expected only blocks {_QWEN_QUANTIZED_BLOCKS} quantized" ) if qwen_model.is_svdquant: a_prefixes = _module_prefixes(keys, ".svdquant_lora_a") b_prefixes = _module_prefixes(keys, ".svdquant_lora_b") pqs_prefixes = _module_prefixes(keys, ".pre_quant_scale") assert a_prefixes, "no promoted svdquant_lora_a keys" # Every promoted linear carries lora_a, lora_b, and pre_quant_scale. The # svdquant_skip_layers (text-stream + modulation linears) are quantized # plain NVFP4: they have weight_scale but no low-rank/pre_quant_scale # tensors, so the promoted and quantized sets differ by exactly them. expected_promoted = { f"transformer_blocks.{block}{suffix}" for block in _QWEN_QUANTIZED_BLOCKS for suffix in _QWEN_SVDQUANT_PROMOTED_SUFFIXES } expected_skipped = { f"transformer_blocks.{block}{suffix}" for block in _QWEN_QUANTIZED_BLOCKS for suffix in _QWEN_SVDQUANT_SKIPPED_SUFFIXES } assert a_prefixes == b_prefixes == pqs_prefixes == expected_promoted assert weight_scale_prefixes == expected_promoted | expected_skipped # Rank-consistent shapes; lora_a=[rank, in], lora_b=[out, rank], rank == --lowrank. for key, tensor in lora_tensors.items(): if key.endswith(".svdquant_lora_a"): assert tensor.shape[0] == _QWEN_LORA_RANK else: assert tensor.shape[1] == _QWEN_LORA_RANK # NVFP4 secondary scales are present. assert any(k.endswith(".weight_scale_2") for k in keys) # config schema (modeled on nvfp4_awq). assert quant_config.get("quant_algo") == "NVFP4_SVD" group = next(iter(quant_config.get("config_groups", {}).values()), {}) assert group.get("lora_rank") == _QWEN_LORA_RANK assert group.get("pre_quant_scale") is True assert group.get("has_zero_point") is False assert quant_config.get("ignore"), "expected excluded modules in 'ignore'" else: # Plain FP8/NVFP4: weight scales present, no SVDQuant tensors. assert weight_scale_prefixes, "no weight_scale in export" assert not any(k.endswith(".svdquant_lora_a") for k in keys) if qwen_model.format_type == "fp4": assert any(k.endswith(".weight_scale_2") for k in keys)