From fc27874a8ea21240bcd21ffd9c23b6ef6fffed0e Mon Sep 17 00:00:00 2001 From: Pham Hong Vinh Date: Tue, 25 Aug 2026 18:40:58 +0000 Subject: [PATCH 1/7] Add data-free SVDQuant quantize-on-load to Nunchaku Lite MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Support `pre_quantized=False` in NunchakuLiteQuantizationConfig: targeted linears of an unquantized checkpoint are quantized at load time with data-free SVDQuant (weight-span smoothing, rank-r SVD low-rank branch, int4/nvfp4 group quantization) and packed directly into the kernel layout SVDQW4A4Linear consumes — no calibration data needed. The math lives in quantizers/nunchaku/data_free.py, which is pure torch and stays importable without the `kernels` package; quantization happens per weight in create_quantized_param so peak memory stays near the quantized model size. Packed outputs are tensor-for-tensor identical to DeepCompressor's Nunchaku W4A4 converter (verified against it on random weights; qweight byte-identical). `awq_w4a16` targets are not supported in this mode and raise. CPU tests validate shapes, round-trip reconstruction error, bias packing, and the quantizer flow; a gated GPU mixin test runs quantize-on-load end to end where kernels are available. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_013NAtDkGmAw2fzbcvjfC79w --- docs/source/en/quantization/nunchaku.md | 37 +++ .../quantizers/nunchaku/data_free.py | 274 ++++++++++++++++++ .../quantizers/nunchaku/nunchaku_quantizer.py | 59 +++- .../quantizers/quantization_config.py | 7 +- tests/models/testing_utils/quantization.py | 21 ++ tests/quantization/nunchaku/__init__.py | 0 tests/quantization/nunchaku/test_data_free.py | 271 +++++++++++++++++ 7 files changed, 666 insertions(+), 3 deletions(-) create mode 100644 src/diffusers/quantizers/nunchaku/data_free.py create mode 100644 tests/quantization/nunchaku/__init__.py create mode 100644 tests/quantization/nunchaku/test_data_free.py diff --git a/docs/source/en/quantization/nunchaku.md b/docs/source/en/quantization/nunchaku.md index d6a07591b6fe..46576c3b8766 100644 --- a/docs/source/en/quantization/nunchaku.md +++ b/docs/source/en/quantization/nunchaku.md @@ -122,6 +122,43 @@ List each module you want to quantize under `svdq_w4a4` or `awq_w4a16`. A module } ``` +## Data-free quantization on load + +Pass `pre_quantized=False` to quantize an *unquantized* checkpoint at load time with data-free SVDQuant — no calibration data is needed. Each targeted linear gets weight-span smoothing, a rank-`r` SVD low-rank branch, and int4 or nvfp4 group quantization of the residual, packed directly into the kernel layout. Only `svdq_w4a4` targets are supported in this mode, and each target's `in_features`/`out_features` must be multiples of 128 (with `rank` a multiple of 16, or 0 to disable the low-rank branch). + +Targets are explicit module paths, so build the list from the model's structure: + +```python +import torch +from diffusers import Flux2Transformer2DModel, NunchakuLiteQuantizationConfig + +model_id = "black-forest-labs/FLUX.2-klein-9B" +with torch.device("meta"): + reference = Flux2Transformer2DModel.from_config( + Flux2Transformer2DModel.load_config(model_id, subfolder="transformer") + ) +targets = [ + name + for name, module in reference.named_modules() + if isinstance(module, torch.nn.Linear) + and name.startswith(("transformer_blocks.", "single_transformer_blocks.")) + and "norm" not in name +] + +transformer = Flux2Transformer2DModel.from_pretrained( + model_id, + subfolder="transformer", + quantization_config=NunchakuLiteQuantizationConfig( + svdq_w4a4={"precision": "nvfp4", "group_size": 16, "rank": 32, "targets": targets}, + pre_quantized=False, + ), + torch_dtype=torch.bfloat16, + device_map="cuda", +) +``` + +Quantization happens per weight as the checkpoint streams in, so peak memory stays near the quantized model size. The result is identical to loading a checkpoint produced offline by a data-free SVDQuant exporter. + ## Fused kernels The original [Nunchaku](https://github.com/nunchaku-ai/nunchaku) engine gets much of its speed from model-specific fused execution paths. It combines the Q, K, and V projections with RMSNorm and RoPE, and uses a fused GELU kernel for the MLP. Nunchaku Lite instead uses the standard Diffusers model with generic quantized linear layers, so it does not include these fusions. diff --git a/src/diffusers/quantizers/nunchaku/data_free.py b/src/diffusers/quantizers/nunchaku/data_free.py new file mode 100644 index 000000000000..fba5bd876e80 --- /dev/null +++ b/src/diffusers/quantizers/nunchaku/data_free.py @@ -0,0 +1,274 @@ +"""Data-free SVDQuant quantization for the Nunchaku Lite backend. + +Quantizes a bf16 linear weight at load time — no calibration data required — +into the exact packed parameter layout consumed by ``SVDQW4A4Linear``: +weight-span smoothing, a rank-``r`` SVD low-rank branch, and int4/nvfp4 group +quantization of the residual. The packing mirrors DeepCompressor's Nunchaku +W4A4 converter, so the produced tensors are indistinguishable from a +pre-quantized checkpoint's. + +This module is pure PyTorch and must stay importable without the ``kernels`` +package (unlike ``.utils``, which fetches the CUDA kernels at import time). +""" + +from __future__ import annotations + +import torch + + +_SMOOTH_EPS = 1e-6 +_FP8_MAX = 448.0 + + +def _ceil_divide(x: int, divisor: int) -> int: + return (x + divisor - 1) // divisor + + +def _pad( + tensor: torch.Tensor, divisor: tuple[int, ...], dim: tuple[int, ...], fill_value: float = 0.0 +) -> torch.Tensor: + shape = list(tensor.shape) + for axis, axis_divisor in zip(dim, divisor): + shape[axis] = _ceil_divide(shape[axis], axis_divisor) * axis_divisor + if shape == list(tensor.shape): + return tensor + result = torch.full(shape, fill_value, dtype=tensor.dtype, device=tensor.device) + result[tuple(slice(0, extent) for extent in tensor.shape)] = tensor + return result + + +def _fp4_e2m1_codebook(device: torch.device, dtype: torch.dtype = torch.float32) -> torch.Tensor: + return torch.tensor( + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0], + dtype=dtype, + device=device, + ) + + +def _fp_quantize(x: torch.Tensor) -> torch.Tensor: + """Quantize values to the nearest FP4 E2M1 codebook index.""" + + codebook = _fp4_e2m1_codebook(x.device, x.dtype) + positive = codebook[:8] + thresholds = (positive[:-1] + positive[1:]) / 2 + codes = torch.bucketize(x.abs(), thresholds, right=False) + negative = x.lt(0) & codes.ne(0) + codes.add_(negative, alpha=8) + codes.masked_fill_(~x.isfinite(), 0) + return codes + + +class _NunchakuWeightPacker: + """Pack-only subset of DeepCompressor's Nunchaku MMA weight packer (4-bit).""" + + def __init__(self, warp_n: int = 128): + self.bits = 4 + self.comp_n = 16 + self.comp_k = 256 // self.bits + self.insn_k = self.comp_k + self.num_lanes = 32 + self.num_k_lanes = 4 + self.num_n_lanes = 8 + self.warp_n = warp_n + self.reg_k = 32 // self.bits + self.reg_n = 1 + self.k_pack_size = self.comp_k // (self.num_k_lanes * self.reg_k) + self.n_pack_size = self.comp_n // (self.num_n_lanes * self.reg_n) + self.mem_k = self.comp_k + self.mem_n = warp_n + self.num_k_packs = self.mem_k // (self.k_pack_size * self.num_k_lanes * self.reg_k) + self.num_n_packs = self.mem_n // (self.n_pack_size * self.num_n_lanes * self.reg_n) + self.num_k_unrolls = 2 + + def pack_weight(self, weight: torch.Tensor) -> torch.Tensor: + weight = _pad(weight, divisor=(self.mem_n, self.mem_k * self.num_k_unrolls), dim=(0, 1)) + n, k = weight.shape + weight = weight.reshape( + n // self.mem_n, + self.num_n_packs, + self.n_pack_size, + self.num_n_lanes, + self.reg_n, + k // self.mem_k, + self.num_k_packs, + self.k_pack_size, + self.num_k_lanes, + self.reg_k, + ) + weight = weight.permute(0, 5, 6, 1, 3, 8, 2, 7, 4, 9).contiguous() + weight = weight.bitwise_and_(0xF) + shift = torch.arange(0, 32, 4, dtype=torch.int32, device=weight.device) + weight = weight.bitwise_left_shift_(shift).sum(dim=-1, dtype=torch.int32) + return weight.view(dtype=torch.int8).view(n, -1) + + def pack_vector(self, vector: torch.Tensor) -> torch.Tensor: + """Pack a per-channel vector (smooth factor, bias) into scale layout.""" + + vector = _pad(vector, divisor=(self.warp_n,), dim=(0,), fill_value=1.0) + n = vector.shape[0] + s_pack_size = min(max(self.warp_n // self.num_lanes, 2), 8) + num_s_lanes = min(self.num_lanes, self.warp_n // s_pack_size) + num_s_packs = self.warp_n // (s_pack_size * num_s_lanes) + vector = vector.reshape(n // self.warp_n, num_s_packs, num_s_lanes // 4, s_pack_size // 2, 4, 2, -1) + vector = vector.permute(0, 6, 1, 2, 4, 3, 5).contiguous() + return vector.view(-1) + + def pack_group_scale(self, scale: torch.Tensor) -> torch.Tensor: + """Pack per-group scales in ``[out, groups]`` layout (int4, group size 64).""" + + scale = _pad( + scale.view(scale.shape[0], 1, -1, 1), divisor=(self.warp_n, self.num_k_unrolls), dim=(0, 2), fill_value=1.0 + ) + n = scale.shape[0] + s_pack_size = min(max(self.warp_n // self.num_lanes, 2), 8) + num_s_lanes = min(self.num_lanes, self.warp_n // s_pack_size) + num_s_packs = self.warp_n // (s_pack_size * num_s_lanes) + scale = scale.reshape(n // self.warp_n, num_s_packs, num_s_lanes // 4, s_pack_size // 2, 4, 2, -1) + scale = scale.permute(0, 6, 1, 2, 4, 3, 5).contiguous() + return scale.view(-1, n) + + def pack_micro_scale(self, scale: torch.Tensor) -> torch.Tensor: + """Pack FP8 per-group scales in ``[out, groups]`` layout (nvfp4, group size 16).""" + + group_fragment = self.insn_k // 16 + scale = _pad( + scale.view(scale.shape[0], 1, -1, 1), divisor=(self.warp_n, group_fragment), dim=(0, 2), fill_value=1.0 + ) + scale = scale.to(dtype=torch.float8_e4m3fn) + n = scale.shape[0] + s_pack_size = min(max(self.warp_n // self.num_lanes, 1), 4) + num_s_lanes = 32 + num_s_packs = _ceil_divide(self.warp_n, s_pack_size * num_s_lanes) + scale = scale.view(n // self.warp_n, num_s_packs, s_pack_size, 4, 8, -1, group_fragment) + scale = scale.permute(0, 5, 1, 4, 3, 2, 6).contiguous() + return scale.view(-1, n) + + def pack_lowrank_weight(self, weight: torch.Tensor, down: bool) -> torch.Tensor: + reg_n, reg_k = 1, 2 + pack_n = self.n_pack_size * self.num_n_lanes * reg_n + pack_k = self.k_pack_size * self.num_k_lanes * reg_k + weight = _pad(weight, divisor=(pack_n, pack_k), dim=(0, 1)) + if down: + r, c = weight.shape + r_packs, c_packs = r // pack_n, c // pack_k + weight = weight.view(r_packs, pack_n, c_packs, pack_k).permute(2, 0, 1, 3) + else: + c, r = weight.shape + c_packs, r_packs = c // pack_n, r // pack_k + weight = weight.view(c_packs, pack_n, r_packs, pack_k).permute(0, 2, 1, 3) + weight = weight.reshape( + c_packs, r_packs, self.n_pack_size, self.num_n_lanes, reg_n, self.k_pack_size, self.num_k_lanes, reg_k + ) + weight = weight.permute(0, 1, 3, 6, 2, 5, 4, 7).contiguous() + return weight.view(c, r) + + +def _check_packable(out_features: int, in_features: int, rank: int, group_size: int) -> None: + if out_features % 128 != 0 or in_features % 128 != 0: + raise ValueError( + "Data-free Nunchaku quantization requires in_features and out_features to be multiples of 128, " + f"got ({out_features}, {in_features})." + ) + if in_features % group_size != 0: + raise ValueError(f"in_features ({in_features}) must be divisible by group_size ({group_size}).") + if rank % 16 != 0: + raise ValueError(f"Low-rank branch rank must be a multiple of 16 (or 0), got {rank}.") + + +def _weight_span_smooth_scale(weight: torch.Tensor) -> torch.Tensor: + """Data-free weight-span smoothing: ``s_j = 1 / absmax(W[:, j]) ** 0.5``. + + The weight is stored multiplied by ``s`` (equalizing per-channel magnitudes) + and the kernel divides the activations by ``s`` at runtime. + """ + + span = weight.abs().amax(dim=0).clamp_min(_SMOOTH_EPS) + scale = 1.0 / span.pow(0.5) + scale = torch.where(torch.isfinite(scale), scale, torch.ones_like(scale)) + return scale.clamp_min(_SMOOTH_EPS) + + +def _group_scales(residual: torch.Tensor, group_size: int, float_point: bool) -> torch.Tensor: + out_features, in_features = residual.shape + groups = in_features // group_size + max_q = 6.0 if float_point else 7.0 + return residual.view(out_features, groups, group_size).abs().amax(dim=2).clamp_min(1e-6) / max_q + + +def quantize_linear_data_free( + weight: torch.Tensor, + *, + precision: str, + group_size: int, + rank: int, + torch_dtype: torch.dtype = torch.bfloat16, +) -> dict[str, torch.Tensor]: + """Quantize one linear weight into ``SVDQW4A4Linear``'s packed parameters. + + Args: + weight: Unquantized weight in ``[out_features, in_features]`` layout. + precision: ``"int4"`` or ``"nvfp4"``. + group_size: Weight quantization group size (64 for int4, 16 for nvfp4). + rank: Low-rank branch rank (multiple of 16, or 0 to disable). + torch_dtype: Floating-point dtype of the produced auxiliary tensors. + + Returns: + Mapping with keys ``qweight``, ``wscales``, ``smooth_factor``, + ``proj_down``, ``proj_up`` and, for nvfp4, ``wcscales`` and ``wtscale``. + """ + + out_features, in_features = weight.shape + _check_packable(out_features, in_features, rank, group_size) + packer = _NunchakuWeightPacker() + weight = weight.to(dtype=torch.float32) + + smooth = _weight_span_smooth_scale(weight) + smoothed = weight * smooth.view(1, -1) + + if rank > 0: + u, s, vh = torch.linalg.svd(smoothed, full_matrices=False) + proj_up = (u[:, :rank] * s[:rank].view(1, -1)).contiguous() + proj_down = vh[:rank, :].contiguous() + residual = smoothed - proj_up @ proj_down + else: + proj_up = smoothed.new_zeros((out_features, 0)) + proj_down = smoothed.new_zeros((0, in_features)) + residual = smoothed + + groups = in_features // group_size + state: dict[str, torch.Tensor] = {} + if precision == "nvfp4": + effective = _group_scales(residual, group_size, float_point=True) + wtscale = (effective.amax() / _FP8_MAX).clamp_min(1e-12) + subscale = (effective / wtscale).clamp(min=0.0, max=_FP8_MAX) + subscale = subscale.to(dtype=torch.float8_e4m3fn).to(dtype=torch.float32) + divisor = (subscale * wtscale).view(out_features, groups, 1) + scaled = residual.view(out_features, groups, group_size) / divisor + codes = _fp_quantize(scaled.reshape(out_features, in_features)).to(torch.int32) + state["wscales"] = packer.pack_micro_scale(subscale) + state["wcscales"] = torch.ones(out_features, dtype=torch_dtype, device=weight.device) + state["wtscale"] = wtscale.view(1).to(dtype=torch_dtype) + elif precision == "int4": + scale = _group_scales(residual, group_size, float_point=False) + scaled = residual.view(out_features, groups, group_size) / scale.view(out_features, groups, 1) + codes = scaled.reshape(out_features, in_features).round_().clamp_(-8, 7).to(torch.int32) + state["wscales"] = packer.pack_group_scale(scale.to(dtype=torch_dtype)) + else: + raise ValueError(f"Unsupported precision for data-free quantization: {precision!r}") + + state["qweight"] = packer.pack_weight(codes) + state["smooth_factor"] = packer.pack_vector(smooth.to(dtype=torch_dtype)) + # The kernel's low-rank branch consumes the unsmoothed input, so fold 1/smooth + # into the down projection; the residual weight stays in smoothed coordinates. + proj_down = proj_down / smooth.view(1, -1) + state["proj_down"] = packer.pack_lowrank_weight(proj_down.to(dtype=torch_dtype), down=True) + state["proj_up"] = packer.pack_lowrank_weight(proj_up.to(dtype=torch_dtype), down=False) + return state + + +def pack_data_free_bias(bias: torch.Tensor, torch_dtype: torch.dtype = torch.bfloat16) -> torch.Tensor: + """Pack a bias vector into the layout ``SVDQW4A4Linear.bias`` expects.""" + + packer = _NunchakuWeightPacker() + packed = packer.pack_vector(bias.to(dtype=torch.float32)) + return packed[: bias.shape[0]].to(dtype=torch_dtype) diff --git a/src/diffusers/quantizers/nunchaku/nunchaku_quantizer.py b/src/diffusers/quantizers/nunchaku/nunchaku_quantizer.py index b8f20d7ddba3..30ae4bfd888b 100644 --- a/src/diffusers/quantizers/nunchaku/nunchaku_quantizer.py +++ b/src/diffusers/quantizers/nunchaku/nunchaku_quantizer.py @@ -6,6 +6,8 @@ if TYPE_CHECKING: + import torch + from ...models.modeling_utils import ModelMixin @@ -19,7 +21,9 @@ class NunchakuLiteQuantizer(DiffusersQuantizer): def __init__(self, quantization_config, **kwargs): super().__init__(quantization_config, **kwargs) self.compute_dtype = quantization_config.compute_dtype - self.pre_quantized = quantization_config.pre_quantized + # Quantize on load when either the loader inferred an unquantized + # checkpoint or the config explicitly requested `pre_quantized=False`. + self.pre_quantized = self.pre_quantized and quantization_config.pre_quantized def validate_environment(self, *args, **kwargs): if not is_kernels_available(): @@ -69,10 +73,61 @@ def _process_model_before_weight_loading( quantization_config = self.quantization_config.to_dict() num_replaced = replace_with_nunchaku_linear(model, quantization_config, self.compute_dtype) - if state_dict is not None: + if self.pre_quantized and state_dict is not None: check_strict_state_dict_match(model, state_dict) logger.info(f"Applied Nunchaku quantization config with {num_replaced} targets.") + def check_if_quantized_param( + self, + model: "ModelMixin", + param_value: "torch.Tensor", + param_name: str, + state_dict: dict[str, Any], + **kwargs, + ) -> bool: + if self.pre_quantized: + return False + from .utils import SVDQW4A4Linear + + module_name, _, tensor_name = param_name.rpartition(".") + if tensor_name not in ("weight", "bias") or not module_name: + return False + try: + module = model.get_submodule(module_name) + except AttributeError: + return False + return isinstance(module, SVDQW4A4Linear) + + def create_quantized_param( + self, + model: "ModelMixin", + param_value: "torch.Tensor", + param_name: str, + target_device: "torch.device", + state_dict: dict[str, Any] | None = None, + unexpected_keys: list[str] | None = None, + **kwargs, + ): + import torch + + from .data_free import pack_data_free_bias, quantize_linear_data_free + + module_name, _, tensor_name = param_name.rpartition(".") + module = model.get_submodule(module_name) + if tensor_name == "bias": + packed_bias = pack_data_free_bias(param_value.to(target_device), torch_dtype=self.compute_dtype) + module._parameters["bias"] = torch.nn.Parameter(packed_bias, requires_grad=False) + return + quantized = quantize_linear_data_free( + param_value.to(target_device), + precision=module.precision, + group_size=module.group_size, + rank=module.rank, + torch_dtype=self.compute_dtype, + ) + for name, tensor in quantized.items(): + module._parameters[name] = torch.nn.Parameter(tensor.to(target_device), requires_grad=False) + def _process_model_after_weight_loading(self, model: "ModelMixin", **kwargs): return model diff --git a/src/diffusers/quantizers/quantization_config.py b/src/diffusers/quantizers/quantization_config.py index ea78b5f7ff53..d8c3bc6bbbc6 100644 --- a/src/diffusers/quantizers/quantization_config.py +++ b/src/diffusers/quantizers/quantization_config.py @@ -492,7 +492,7 @@ def __init__( if not isinstance(compute_dtype, torch.dtype): raise ValueError("Nunchaku compute_dtype must be a string or a torch.dtype.") self.compute_dtype = compute_dtype - self.pre_quantized = True + self.pre_quantized = kwargs.pop("pre_quantized", True) self.svdq_w4a4 = svdq_w4a4 self.awq_w4a16 = awq_w4a16 @@ -503,6 +503,11 @@ def post_init(self): raise ValueError( "Nunchaku compact quantization config must include `svdq_w4a4.targets` or `awq_w4a16.targets`." ) + if not self.pre_quantized and self.awq_w4a16 is not None: + raise NotImplementedError( + "Data-free quantization (`pre_quantized=False`) only supports `svdq_w4a4` targets; " + "remove the `awq_w4a16` section or load a pre-quantized checkpoint." + ) for op, raw in (("svdq_w4a4", self.svdq_w4a4), ("awq_w4a16", self.awq_w4a16)): if raw is None: diff --git a/tests/models/testing_utils/quantization.py b/tests/models/testing_utils/quantization.py index 918126fe3f13..589ef81c5fba 100644 --- a/tests/models/testing_utils/quantization.py +++ b/tests/models/testing_utils/quantization.py @@ -1762,6 +1762,27 @@ def _test_quantized_layers(self, config_kwargs): def test_nunchaku_lite_quantized_layers(self): self._test_quantized_layers(self.config_dict) + def test_nunchaku_lite_data_free_quantization(self): + """Quantize an unquantized checkpoint on load (`pre_quantized=False`) and run a forward pass.""" + + unquantized_path = getattr(self, "unquantized_model_name_or_path", None) + data_free_config = getattr(self, "data_free_config_dict", None) + if unquantized_path is None or data_free_config is None: + pytest.skip("Data-free quantization attributes are not configured for this model.") + + kwargs = getattr(self, "pretrained_model_kwargs", {}).copy() + kwargs["quantization_config"] = NunchakuLiteQuantizationConfig(**data_free_config, pre_quantized=False) + model = self.model_class.from_pretrained(unquantized_path, **kwargs) + + num_quantized_layers = sum(1 for _, module in model.named_modules() if self._is_module_quantized(module)) + expected = len(data_free_config["svdq_w4a4"]["targets"]) + assert num_quantized_layers == expected, ( + f"Data-free quantization replaced {num_quantized_layers} layers, expected {expected}." + ) + + with torch.no_grad(): + model(**self.get_dummy_inputs()) + @pytest.mark.skipif(not is_kernels_available(), reason="`kernels` is not available.") @require_accelerate diff --git a/tests/quantization/nunchaku/__init__.py b/tests/quantization/nunchaku/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/quantization/nunchaku/test_data_free.py b/tests/quantization/nunchaku/test_data_free.py new file mode 100644 index 000000000000..9d2da4985754 --- /dev/null +++ b/tests/quantization/nunchaku/test_data_free.py @@ -0,0 +1,271 @@ +# coding=utf-8 +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# 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. + +"""CPU-only tests for data-free Nunchaku SVDQuant quantization. + +These tests intentionally avoid importing ``diffusers.quantizers.nunchaku.utils`` +(which requires the ``kernels`` package and a CUDA GPU); the packed layouts are +validated against pure-torch reference unpackers ported from DeepCompressor. +""" + +import pytest +import torch + +from diffusers import NunchakuLiteQuantizationConfig +from diffusers.quantizers.nunchaku.data_free import ( + _NunchakuWeightPacker, + pack_data_free_bias, + quantize_linear_data_free, +) + + +# --------------------------------------------------------------------------- +# Reference unpackers (ported from DeepCompressor's Nunchaku converter). +# --------------------------------------------------------------------------- + + +def _ceil_divide(x, divisor): + return (x + divisor - 1) // divisor + + +def _unpack_weight(packed, rows, columns): + p = _NunchakuWeightPacker() + padded_rows = _ceil_divide(rows, p.mem_n) * p.mem_n + padded_columns = _ceil_divide(columns, p.mem_k * p.num_k_unrolls) * p.mem_k * p.num_k_unrolls + unpacked = packed.contiguous().view(torch.int32) + unpacked = unpacked.view( + padded_rows // p.mem_n, + padded_columns // p.mem_k, + p.num_k_packs, + p.num_n_packs, + p.num_n_lanes, + p.num_k_lanes, + p.n_pack_size, + p.k_pack_size, + p.reg_n, + ) + shift = torch.arange(0, 32, 4, dtype=torch.int32) + unpacked = unpacked.unsqueeze(-1).bitwise_right_shift(shift).bitwise_and(0xF) + unpacked = torch.where(unpacked >= 8, unpacked - 16, unpacked) + unpacked = unpacked.permute(0, 3, 6, 4, 8, 1, 2, 7, 5, 9).contiguous() + return unpacked.view(padded_rows, padded_columns)[:rows, :columns] + + +def _unpack_vector(packed, rows): + p = _NunchakuWeightPacker() + padded_rows = _ceil_divide(rows, p.warp_n) * p.warp_n + s_pack_size = min(max(p.warp_n // p.num_lanes, 2), 8) + num_s_lanes = min(p.num_lanes, p.warp_n // s_pack_size) + num_s_packs = p.warp_n // (s_pack_size * num_s_lanes) + unpacked = packed.contiguous().view( + padded_rows // p.warp_n, 1, num_s_packs, num_s_lanes // 4, 4, s_pack_size // 2, 2 + ) + unpacked = unpacked.permute(0, 2, 3, 5, 4, 6, 1).contiguous() + return unpacked.view(padded_rows)[:rows] + + +def _unpack_group_scale(packed, rows, groups): + p = _NunchakuWeightPacker() + padded_rows = _ceil_divide(rows, p.warp_n) * p.warp_n + padded_groups = _ceil_divide(groups, p.num_k_unrolls) * p.num_k_unrolls + s_pack_size = min(max(p.warp_n // p.num_lanes, 2), 8) + num_s_lanes = min(p.num_lanes, p.warp_n // s_pack_size) + num_s_packs = p.warp_n // (s_pack_size * num_s_lanes) + unpacked = packed.contiguous().view( + padded_rows // p.warp_n, padded_groups, num_s_packs, num_s_lanes // 4, 4, s_pack_size // 2, 2 + ) + unpacked = unpacked.permute(0, 2, 3, 5, 4, 6, 1).contiguous() + return unpacked.view(padded_rows, padded_groups)[:rows, :groups] + + +def _unpack_micro_scale(packed, rows, groups): + p = _NunchakuWeightPacker() + padded_rows = _ceil_divide(rows, p.warp_n) * p.warp_n + group_fragment = p.insn_k // 16 + padded_groups = _ceil_divide(groups, group_fragment) * group_fragment + s_pack_size = min(max(p.warp_n // p.num_lanes, 1), 4) + num_s_packs = _ceil_divide(p.warp_n, s_pack_size * 32) + unpacked = packed.contiguous().view( + padded_rows // p.warp_n, padded_groups // group_fragment, num_s_packs, 8, 4, s_pack_size, group_fragment + ) + unpacked = unpacked.permute(0, 2, 5, 4, 3, 1, 6).contiguous() + return unpacked.view(padded_rows, padded_groups)[:rows, :groups] + + +def _unpack_lowrank(packed, down, rows, columns): + p = _NunchakuWeightPacker() + reg_n, reg_k = 1, 2 + pack_n = p.n_pack_size * p.num_n_lanes * reg_n + pack_k = p.k_pack_size * p.num_k_lanes * reg_k + padded_rows = _ceil_divide(rows, pack_n) * pack_n + padded_columns = _ceil_divide(columns, pack_k) * pack_k + if down: + r, c = padded_rows, padded_columns + r_packs, c_packs = r // pack_n, c // pack_k + else: + c, r = padded_rows, padded_columns + c_packs, r_packs = c // pack_n, r // pack_k + unpacked = packed.contiguous().view( + c_packs, r_packs, p.num_n_lanes, p.num_k_lanes, p.n_pack_size, p.k_pack_size, reg_n, reg_k + ) + unpacked = unpacked.permute(0, 1, 4, 2, 6, 5, 3, 7).contiguous() + unpacked = unpacked.view(c_packs, r_packs, pack_n, pack_k) + if down: + unpacked = unpacked.permute(1, 2, 0, 3).contiguous().view(r, c) + else: + unpacked = unpacked.permute(0, 2, 1, 3).contiguous().view(c, r) + return unpacked[:rows, :columns] + + +def _fp4_codebook(): + return torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0]) + + +def _reconstruct(state, out_features, in_features, group_size, rank, precision): + """Rebuild the original weight from the packed data-free state.""" + + codes = _unpack_weight(state["qweight"], out_features, in_features).float() + groups = in_features // group_size + if precision == "nvfp4": + values = _fp4_codebook()[codes.long() & 0xF] + wscales = _unpack_micro_scale(state["wscales"].view(torch.float8_e4m3fn), out_features, groups).float() + scale = wscales * state["wtscale"].float() + else: + values = codes + wscales = _unpack_group_scale(state["wscales"], out_features, groups).float() + scale = wscales + residual = values.view(out_features, groups, group_size) * scale.view(out_features, groups, 1) + residual = residual.view(out_features, in_features) + smooth = _unpack_vector(state["smooth_factor"], in_features).float() + down = _unpack_lowrank(state["proj_down"], down=True, rows=rank, columns=in_features).float() + up = _unpack_lowrank(state["proj_up"], down=False, rows=out_features, columns=rank).float() + # Residual is in smoothed coordinates; the low-rank branch already absorbed 1/smooth. + return residual / smooth.view(1, -1) + up @ down + + +OUT_FEATURES, IN_FEATURES = 256, 384 + + +@pytest.mark.parametrize("precision,group_size", [("int4", 64), ("nvfp4", 16)]) +def test_data_free_state_shapes_and_dtypes(precision, group_size): + weight = torch.randn(OUT_FEATURES, IN_FEATURES) + state = quantize_linear_data_free(weight, precision=precision, group_size=group_size, rank=32) + + assert state["qweight"].shape == (OUT_FEATURES, IN_FEATURES // 2) + assert state["qweight"].dtype == torch.int8 + assert state["smooth_factor"].shape == (IN_FEATURES,) + assert state["proj_down"].shape == (IN_FEATURES, 32) + assert state["proj_up"].shape == (OUT_FEATURES, 32) + assert state["wscales"].shape == (IN_FEATURES // group_size, OUT_FEATURES) + if precision == "nvfp4": + assert state["wscales"].dtype == torch.float8_e4m3fn + assert state["wcscales"].shape == (OUT_FEATURES,) + assert torch.all(state["wcscales"].float() == 1.0) + assert state["wtscale"].shape == (1,) + else: + assert state["wscales"].dtype == torch.bfloat16 + assert "wcscales" not in state + assert "wtscale" not in state + for tensor in (state["smooth_factor"], state["proj_down"], state["proj_up"]): + assert tensor.dtype == torch.bfloat16 + + +@pytest.mark.parametrize("precision,group_size", [("int4", 64), ("nvfp4", 16)]) +def test_data_free_round_trip_error_bounded(precision, group_size): + torch.manual_seed(0) + weight = torch.randn(OUT_FEATURES, IN_FEATURES) + + state = quantize_linear_data_free(weight, precision=precision, group_size=group_size, rank=32) + reconstructed = _reconstruct(state, OUT_FEATURES, IN_FEATURES, group_size, 32, precision) + error = (reconstructed - weight).norm() / weight.norm() + + state_rank0 = quantize_linear_data_free(weight, precision=precision, group_size=group_size, rank=0) + reconstructed_rank0 = _reconstruct(state_rank0, OUT_FEATURES, IN_FEATURES, group_size, 0, precision) + error_rank0 = (reconstructed_rank0 - weight).norm() / weight.norm() + + assert error < 0.15 + assert error < error_rank0 + + +def test_data_free_bias_round_trip(): + torch.manual_seed(0) + bias = torch.randn(OUT_FEATURES) + packed = pack_data_free_bias(bias) + assert packed.shape == (OUT_FEATURES,) + assert packed.dtype == torch.bfloat16 + assert torch.allclose(_unpack_vector(packed, OUT_FEATURES).float(), bias, atol=1e-2, rtol=1e-2) + + +def test_data_free_rejects_unsupported_dimensions(): + with pytest.raises(ValueError, match="multiples of 128"): + quantize_linear_data_free(torch.randn(100, 384), precision="int4", group_size=64, rank=32) + with pytest.raises(ValueError, match="multiple of 16"): + quantize_linear_data_free(torch.randn(256, 384), precision="int4", group_size=64, rank=24) + with pytest.raises(ValueError, match="Unsupported precision"): + quantize_linear_data_free(torch.randn(256, 384), precision="fp8", group_size=64, rank=32) + + +def test_config_accepts_pre_quantized_flag(): + config = NunchakuLiteQuantizationConfig( + svdq_w4a4={"precision": "nvfp4", "group_size": 16, "rank": 32, "targets": ["proj"]} + ) + assert config.pre_quantized is True + + config = NunchakuLiteQuantizationConfig( + svdq_w4a4={"precision": "nvfp4", "group_size": 16, "rank": 32, "targets": ["proj"]}, + pre_quantized=False, + ) + assert config.pre_quantized is False + + +def test_config_rejects_data_free_awq(): + with pytest.raises(NotImplementedError, match="svdq_w4a4"): + NunchakuLiteQuantizationConfig( + awq_w4a16={"precision": "int4", "group_size": 64, "targets": ["proj"]}, + pre_quantized=False, + ) + + +def test_quantizer_create_quantized_param_fills_module(): + from diffusers.quantizers.nunchaku.nunchaku_quantizer import NunchakuLiteQuantizer + + config = NunchakuLiteQuantizationConfig( + svdq_w4a4={"precision": "nvfp4", "group_size": 16, "rank": 32, "targets": ["proj"]}, + pre_quantized=False, + ) + quantizer = NunchakuLiteQuantizer(config, pre_quantized=False) + assert quantizer.pre_quantized is False + + class StubQuantizedLinear(torch.nn.Module): + precision = "nvfp4" + group_size = 16 + rank = 32 + + class StubModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.proj = StubQuantizedLinear() + + model = StubModel() + weight = torch.randn(OUT_FEATURES, IN_FEATURES) + quantizer.create_quantized_param(model, weight, "proj.weight", torch.device("cpu")) + quantizer.create_quantized_param(model, torch.randn(OUT_FEATURES), "proj.bias", torch.device("cpu")) + + parameters = dict(model.proj.named_parameters()) + for name in ("qweight", "wscales", "wcscales", "wtscale", "smooth_factor", "proj_down", "proj_up", "bias"): + assert name in parameters, f"missing quantized parameter {name}" + assert not parameters[name].requires_grad + assert parameters["qweight"].shape == (OUT_FEATURES, IN_FEATURES // 2) + assert parameters["bias"].shape == (OUT_FEATURES,) From ebfc2c96cbcfd3dc719b4a131306aee87b26f08e Mon Sep 17 00:00:00 2001 From: Pham Hong Vinh Date: Tue, 25 Aug 2026 18:46:21 +0000 Subject: [PATCH 2/7] Infer data-free quantization targets automatically When `pre_quantized=False` and `svdq_w4a4.targets` is omitted, the quantizer now infers targets from the model at load time: every nn.Linear whose dimensions satisfy the Nunchaku packing constraints is selected, minus modules matched by the new `modules_to_not_convert` config option or listed in the model's `_keep_in_fp32_modules`. Explicit target lists keep working. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_013NAtDkGmAw2fzbcvjfC79w --- docs/source/en/quantization/nunchaku.md | 22 ++---- .../quantizers/nunchaku/data_free.py | 33 +++++++++ .../quantizers/nunchaku/nunchaku_quantizer.py | 11 +++ .../quantizers/quantization_config.py | 17 +++-- tests/quantization/nunchaku/test_data_free.py | 68 +++++++++++++++++++ 5 files changed, 129 insertions(+), 22 deletions(-) diff --git a/docs/source/en/quantization/nunchaku.md b/docs/source/en/quantization/nunchaku.md index 46576c3b8766..37a1c132f543 100644 --- a/docs/source/en/quantization/nunchaku.md +++ b/docs/source/en/quantization/nunchaku.md @@ -126,38 +126,26 @@ List each module you want to quantize under `svdq_w4a4` or `awq_w4a16`. A module Pass `pre_quantized=False` to quantize an *unquantized* checkpoint at load time with data-free SVDQuant — no calibration data is needed. Each targeted linear gets weight-span smoothing, a rank-`r` SVD low-rank branch, and int4 or nvfp4 group quantization of the residual, packed directly into the kernel layout. Only `svdq_w4a4` targets are supported in this mode, and each target's `in_features`/`out_features` must be multiples of 128 (with `rank` a multiple of 16, or 0 to disable the low-rank branch). -Targets are explicit module paths, so build the list from the model's structure: +When `targets` is omitted, eligible targets are inferred automatically from the model: every `nn.Linear` whose dimensions satisfy the packing constraints is quantized, except modules matched by `modules_to_not_convert` (substring match) or listed in the model's `_keep_in_fp32_modules`. Precision-critical modules such as embedders, final projections, and modulation layers are good candidates to exclude: ```python import torch from diffusers import Flux2Transformer2DModel, NunchakuLiteQuantizationConfig -model_id = "black-forest-labs/FLUX.2-klein-9B" -with torch.device("meta"): - reference = Flux2Transformer2DModel.from_config( - Flux2Transformer2DModel.load_config(model_id, subfolder="transformer") - ) -targets = [ - name - for name, module in reference.named_modules() - if isinstance(module, torch.nn.Linear) - and name.startswith(("transformer_blocks.", "single_transformer_blocks.")) - and "norm" not in name -] - transformer = Flux2Transformer2DModel.from_pretrained( - model_id, + "black-forest-labs/FLUX.2-klein-9B", subfolder="transformer", quantization_config=NunchakuLiteQuantizationConfig( - svdq_w4a4={"precision": "nvfp4", "group_size": 16, "rank": 32, "targets": targets}, + svdq_w4a4={"precision": "nvfp4", "group_size": 16, "rank": 32}, pre_quantized=False, + modules_to_not_convert=["context_embedder", "proj_out", "norm", "modulation"], ), torch_dtype=torch.bfloat16, device_map="cuda", ) ``` -Quantization happens per weight as the checkpoint streams in, so peak memory stays near the quantized model size. The result is identical to loading a checkpoint produced offline by a data-free SVDQuant exporter. +An explicit `targets` list is still accepted for full control. Quantization happens per weight as the checkpoint streams in, so peak memory stays near the quantized model size. The result is identical to loading a checkpoint produced offline by a data-free SVDQuant exporter. ## Fused kernels diff --git a/src/diffusers/quantizers/nunchaku/data_free.py b/src/diffusers/quantizers/nunchaku/data_free.py index fba5bd876e80..c5af7039776f 100644 --- a/src/diffusers/quantizers/nunchaku/data_free.py +++ b/src/diffusers/quantizers/nunchaku/data_free.py @@ -163,6 +163,39 @@ def pack_lowrank_weight(self, weight: torch.Tensor, down: bool) -> torch.Tensor: return weight.view(c, r) +def infer_data_free_targets( + model: "torch.nn.Module", + *, + group_size: int, + modules_to_not_convert: tuple[str, ...] | list[str] = (), +) -> list[str]: + """Infer quantization targets for data-free mode from a model's structure. + + Every ``nn.Linear`` whose dimensions fit the Nunchaku packing constraints + (``in_features``/``out_features`` multiples of 128 and ``in_features`` + divisible by ``group_size``) is selected, unless its module path contains + one of the ``modules_to_not_convert`` substrings or the model lists it in + ``_keep_in_fp32_modules``. + """ + + exclude = list(modules_to_not_convert) + list(getattr(model, "_keep_in_fp32_modules", None) or []) + targets = [] + for name, module in model.named_modules(): + if not isinstance(module, torch.nn.Linear): + continue + if any(pattern in name for pattern in exclude): + continue + if module.out_features % 128 or module.in_features % 128 or module.in_features % group_size: + continue + targets.append(name) + if not targets: + raise ValueError( + "Could not infer any data-free quantization targets: no nn.Linear module satisfies the " + "Nunchaku packing constraints (in/out features multiples of 128) outside the excluded modules." + ) + return targets + + def _check_packable(out_features: int, in_features: int, rank: int, group_size: int) -> None: if out_features % 128 != 0 or in_features % 128 != 0: raise ValueError( diff --git a/src/diffusers/quantizers/nunchaku/nunchaku_quantizer.py b/src/diffusers/quantizers/nunchaku/nunchaku_quantizer.py index 30ae4bfd888b..21367573bb4d 100644 --- a/src/diffusers/quantizers/nunchaku/nunchaku_quantizer.py +++ b/src/diffusers/quantizers/nunchaku/nunchaku_quantizer.py @@ -70,6 +70,17 @@ def _process_model_before_weight_loading( ): from .utils import check_strict_state_dict_match, replace_with_nunchaku_linear + svdq_config = self.quantization_config.svdq_w4a4 + if not self.pre_quantized and svdq_config is not None and svdq_config.get("targets") is None: + from .data_free import infer_data_free_targets + + svdq_config["targets"] = infer_data_free_targets( + model, + group_size=svdq_config["group_size"], + modules_to_not_convert=self.quantization_config.modules_to_not_convert or (), + ) + logger.info(f"Inferred {len(svdq_config['targets'])} data-free quantization targets.") + quantization_config = self.quantization_config.to_dict() num_replaced = replace_with_nunchaku_linear(model, quantization_config, self.compute_dtype) diff --git a/src/diffusers/quantizers/quantization_config.py b/src/diffusers/quantizers/quantization_config.py index d8c3bc6bbbc6..c41c4f638c16 100644 --- a/src/diffusers/quantizers/quantization_config.py +++ b/src/diffusers/quantizers/quantization_config.py @@ -493,6 +493,7 @@ def __init__( raise ValueError("Nunchaku compute_dtype must be a string or a torch.dtype.") self.compute_dtype = compute_dtype self.pre_quantized = kwargs.pop("pre_quantized", True) + self.modules_to_not_convert = kwargs.pop("modules_to_not_convert", None) self.svdq_w4a4 = svdq_w4a4 self.awq_w4a16 = awq_w4a16 @@ -515,8 +516,13 @@ def post_init(self): if not isinstance(raw, dict): raise ValueError(f"Nunchaku compact config section {op!r} must be a JSON object.") + # In data-free mode (`pre_quantized=False`) `targets` may be omitted; + # the quantizer infers them from the model at load time. + targets_optional = op == "svdq_w4a4" and not self.pre_quantized for key, expected_type in (("precision", str), ("group_size", int), ("targets", list)): if key not in raw: + if key == "targets" and targets_optional: + continue raise ValueError(f"Nunchaku compact config section {op!r} is missing required field {key!r}.") if not isinstance(raw[key], expected_type): raise ValueError( @@ -525,15 +531,16 @@ def post_init(self): precision = raw["precision"] group_size = raw["group_size"] - targets = raw["targets"] + targets = raw.get("targets") if precision not in ("int4", "nvfp4"): raise ValueError(f"Unsupported Nunchaku precision {precision!r} for {op!r}.") if group_size <= 0: raise ValueError(f"Nunchaku compact config section {op!r} must have positive group_size.") - if not targets: - raise ValueError(f"Nunchaku compact config section {op!r} must contain at least one target.") - if not all(isinstance(target, str) for target in targets): - raise ValueError(f"Nunchaku compact config section {op!r} targets must be strings.") + if targets is not None or not targets_optional: + if not targets: + raise ValueError(f"Nunchaku compact config section {op!r} must contain at least one target.") + if not all(isinstance(target, str) for target in targets): + raise ValueError(f"Nunchaku compact config section {op!r} targets must be strings.") if op == "svdq_w4a4": if "rank" not in raw: diff --git a/tests/quantization/nunchaku/test_data_free.py b/tests/quantization/nunchaku/test_data_free.py index 9d2da4985754..17c9f4e882d4 100644 --- a/tests/quantization/nunchaku/test_data_free.py +++ b/tests/quantization/nunchaku/test_data_free.py @@ -269,3 +269,71 @@ def __init__(self): assert not parameters[name].requires_grad assert parameters["qweight"].shape == (OUT_FEATURES, IN_FEATURES // 2) assert parameters["bias"].shape == (OUT_FEATURES,) + + +class _InferenceToyModel(torch.nn.Module): + _keep_in_fp32_modules = ["frozen"] + + def __init__(self): + super().__init__() + self.blocks = torch.nn.ModuleList( + [torch.nn.Sequential(torch.nn.Linear(IN_FEATURES, OUT_FEATURES)) for _ in range(2)] + ) + self.embedder = torch.nn.Linear(IN_FEATURES, OUT_FEATURES) + self.frozen = torch.nn.Linear(IN_FEATURES, OUT_FEATURES) + self.odd_shape = torch.nn.Linear(100, OUT_FEATURES) + self.norm = torch.nn.LayerNorm(OUT_FEATURES) + + +def test_infer_data_free_targets(): + from diffusers.quantizers.nunchaku.data_free import infer_data_free_targets + + model = _InferenceToyModel() + targets = infer_data_free_targets(model, group_size=16) + # `frozen` is excluded via _keep_in_fp32_modules; `odd_shape` fails the 128-multiple constraint. + assert targets == ["blocks.0.0", "blocks.1.0", "embedder"] + + targets = infer_data_free_targets(model, group_size=16, modules_to_not_convert=["embedder"]) + assert targets == ["blocks.0.0", "blocks.1.0"] + + with pytest.raises(ValueError, match="Could not infer"): + infer_data_free_targets(model, group_size=16, modules_to_not_convert=["blocks", "embedder"]) + + +def test_quantizer_infers_targets_when_omitted(monkeypatch): + import sys + import types + + from diffusers.quantizers.nunchaku.nunchaku_quantizer import NunchakuLiteQuantizer + + config = NunchakuLiteQuantizationConfig( + svdq_w4a4={"precision": "nvfp4", "group_size": 16, "rank": 32}, + pre_quantized=False, + modules_to_not_convert=["embedder"], + ) + quantizer = NunchakuLiteQuantizer(config, pre_quantized=False) + model = _InferenceToyModel() + + # Stub out `.utils` (its import fetches the CUDA kernels) so only the + # target-inference part of _process_model_before_weight_loading runs. + stub = types.ModuleType("diffusers.quantizers.nunchaku.utils") + stub.replace_with_nunchaku_linear = lambda target_model, quantization_config, compute_dtype: len( + quantization_config["svdq_w4a4"]["targets"] + ) + stub.check_strict_state_dict_match = None + monkeypatch.setitem(sys.modules, "diffusers.quantizers.nunchaku.utils", stub) + + quantizer._process_model_before_weight_loading(model) + + assert config.svdq_w4a4["targets"] == ["blocks.0.0", "blocks.1.0"] + + +def test_config_targets_optional_only_for_data_free(): + config = NunchakuLiteQuantizationConfig( + svdq_w4a4={"precision": "nvfp4", "group_size": 16, "rank": 32}, + pre_quantized=False, + ) + assert config.svdq_w4a4.get("targets") is None + + with pytest.raises(ValueError, match="missing required field 'targets'"): + NunchakuLiteQuantizationConfig(svdq_w4a4={"precision": "nvfp4", "group_size": 16, "rank": 32}) From 26bf67ede3a7bcc0820c0165a155bcb06fcd2735 Mon Sep 17 00:00:00 2001 From: Pham Hong Vinh Date: Tue, 25 Aug 2026 18:54:36 +0000 Subject: [PATCH 3/7] Infer data-free exclusions structurally Auto-inference no longer needs a hand-written modules_to_not_convert list: targets are restricted to the model's repeated transformer-block stacks (identical-class nn.ModuleLists), which structurally excludes embedders, final projections, and modulation heads, and adaLN-style linears inside blocks are skipped via default ("norm", "modulation") name patterns. An explicit modules_to_not_convert replaces the default patterns. For FLUX.2-klein-9B the zero-config inferred target set matches the curated list exactly (144 targets). Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_013NAtDkGmAw2fzbcvjfC79w --- docs/source/en/quantization/nunchaku.md | 5 +-- .../quantizers/nunchaku/data_free.py | 40 +++++++++++++++++-- tests/quantization/nunchaku/test_data_free.py | 37 +++++++++++------ 3 files changed, 62 insertions(+), 20 deletions(-) diff --git a/docs/source/en/quantization/nunchaku.md b/docs/source/en/quantization/nunchaku.md index 37a1c132f543..e303dba972b9 100644 --- a/docs/source/en/quantization/nunchaku.md +++ b/docs/source/en/quantization/nunchaku.md @@ -126,7 +126,7 @@ List each module you want to quantize under `svdq_w4a4` or `awq_w4a16`. A module Pass `pre_quantized=False` to quantize an *unquantized* checkpoint at load time with data-free SVDQuant — no calibration data is needed. Each targeted linear gets weight-span smoothing, a rank-`r` SVD low-rank branch, and int4 or nvfp4 group quantization of the residual, packed directly into the kernel layout. Only `svdq_w4a4` targets are supported in this mode, and each target's `in_features`/`out_features` must be multiples of 128 (with `rank` a multiple of 16, or 0 to disable the low-rank branch). -When `targets` is omitted, eligible targets are inferred automatically from the model: every `nn.Linear` whose dimensions satisfy the packing constraints is quantized, except modules matched by `modules_to_not_convert` (substring match) or listed in the model's `_keep_in_fp32_modules`. Precision-critical modules such as embedders, final projections, and modulation layers are good candidates to exclude: +When `targets` is omitted, eligible targets are inferred automatically from the model's structure: quantization is restricted to the repeated transformer-block stacks (so embedders, final projections, and modulation heads outside the stacks stay unquantized), adaLN-style linears are skipped via the default `("norm", "modulation")` name patterns, and every remaining `nn.Linear` satisfying the packing constraints is selected. The model's `_keep_in_fp32_modules` is always honored. No configuration is needed for typical DiTs: ```python import torch @@ -138,14 +138,13 @@ transformer = Flux2Transformer2DModel.from_pretrained( quantization_config=NunchakuLiteQuantizationConfig( svdq_w4a4={"precision": "nvfp4", "group_size": 16, "rank": 32}, pre_quantized=False, - modules_to_not_convert=["context_embedder", "proj_out", "norm", "modulation"], ), torch_dtype=torch.bfloat16, device_map="cuda", ) ``` -An explicit `targets` list is still accepted for full control. Quantization happens per weight as the checkpoint streams in, so peak memory stays near the quantized model size. The result is identical to loading a checkpoint produced offline by a data-free SVDQuant exporter. +Pass `modules_to_not_convert` (substring match; it replaces the default name patterns) to exclude further modules, or an explicit `targets` list for full control. Quantization happens per weight as the checkpoint streams in, so peak memory stays near the quantized model size. The result is identical to loading a checkpoint produced offline by a data-free SVDQuant exporter. ## Fused kernels diff --git a/src/diffusers/quantizers/nunchaku/data_free.py b/src/diffusers/quantizers/nunchaku/data_free.py index c5af7039776f..4a1f00703b06 100644 --- a/src/diffusers/quantizers/nunchaku/data_free.py +++ b/src/diffusers/quantizers/nunchaku/data_free.py @@ -163,26 +163,58 @@ def pack_lowrank_weight(self, weight: torch.Tensor, down: bool) -> torch.Tensor: return weight.view(c, r) +# adaLN-style linears (feature-wise modulation) are precision-critical and are +# consistently named after their norm across diffusers models. +_DEFAULT_EXCLUDE_PATTERNS = ("norm", "modulation") + + +def _repeated_block_prefixes(model: "torch.nn.Module") -> list[str]: + """Return prefixes of ``nn.ModuleList`` stacks of repeated block classes. + + Diffusion transformers keep their compute-heavy linears inside stacks of + identical blocks; peripheral modules (embedders, final projections, + modulation heads) live outside them and should stay unquantized. + """ + + prefixes = [] + for name, module in model.named_modules(): + if not isinstance(module, torch.nn.ModuleList) or len(module) < 2: + continue + if len({type(child) for child in module}) != 1: + continue + prefixes.append(f"{name}.") + return prefixes + + def infer_data_free_targets( model: "torch.nn.Module", *, group_size: int, - modules_to_not_convert: tuple[str, ...] | list[str] = (), + modules_to_not_convert: tuple[str, ...] | list[str] | None = None, ) -> list[str]: """Infer quantization targets for data-free mode from a model's structure. Every ``nn.Linear`` whose dimensions fit the Nunchaku packing constraints (``in_features``/``out_features`` multiples of 128 and ``in_features`` - divisible by ``group_size``) is selected, unless its module path contains - one of the ``modules_to_not_convert`` substrings or the model lists it in - ``_keep_in_fp32_modules``. + divisible by ``group_size``) is selected, restricted to the repeated block + stacks of the model (when it has any) so that peripheral modules such as + embedders and final projections stay unquantized. Modules whose path + contains a ``modules_to_not_convert`` substring — defaulting to + ``("norm", "modulation")`` to skip adaLN-style linears — or matches the + model's ``_keep_in_fp32_modules`` are excluded. Pass an explicit (possibly + empty) ``modules_to_not_convert`` list to replace the default patterns. """ + if modules_to_not_convert is None: + modules_to_not_convert = _DEFAULT_EXCLUDE_PATTERNS exclude = list(modules_to_not_convert) + list(getattr(model, "_keep_in_fp32_modules", None) or []) + stack_prefixes = _repeated_block_prefixes(model) targets = [] for name, module in model.named_modules(): if not isinstance(module, torch.nn.Linear): continue + if stack_prefixes and not any(name.startswith(prefix) for prefix in stack_prefixes): + continue if any(pattern in name for pattern in exclude): continue if module.out_features % 128 or module.in_features % 128 or module.in_features % group_size: diff --git a/tests/quantization/nunchaku/test_data_free.py b/tests/quantization/nunchaku/test_data_free.py index 17c9f4e882d4..d0b4b52b7997 100644 --- a/tests/quantization/nunchaku/test_data_free.py +++ b/tests/quantization/nunchaku/test_data_free.py @@ -271,33 +271,44 @@ def __init__(self): assert parameters["bias"].shape == (OUT_FEATURES,) +class _ToyBlock(torch.nn.Module): + def __init__(self): + super().__init__() + self.proj = torch.nn.Linear(IN_FEATURES, OUT_FEATURES) + self.norm_linear = torch.nn.Linear(IN_FEATURES, OUT_FEATURES) # adaLN-style + self.frozen = torch.nn.Linear(IN_FEATURES, OUT_FEATURES) + self.odd_shape = torch.nn.Linear(100, OUT_FEATURES) + + class _InferenceToyModel(torch.nn.Module): _keep_in_fp32_modules = ["frozen"] def __init__(self): super().__init__() - self.blocks = torch.nn.ModuleList( - [torch.nn.Sequential(torch.nn.Linear(IN_FEATURES, OUT_FEATURES)) for _ in range(2)] - ) + self.blocks = torch.nn.ModuleList([_ToyBlock() for _ in range(2)]) self.embedder = torch.nn.Linear(IN_FEATURES, OUT_FEATURES) - self.frozen = torch.nn.Linear(IN_FEATURES, OUT_FEATURES) - self.odd_shape = torch.nn.Linear(100, OUT_FEATURES) - self.norm = torch.nn.LayerNorm(OUT_FEATURES) + self.proj_out = torch.nn.Linear(OUT_FEATURES, IN_FEATURES) def test_infer_data_free_targets(): from diffusers.quantizers.nunchaku.data_free import infer_data_free_targets model = _InferenceToyModel() + # Default: restricted to the repeated `blocks` stack (embedder/proj_out are + # outside), minus adaLN-style names ("norm"), _keep_in_fp32_modules, and + # dimension-ineligible layers. targets = infer_data_free_targets(model, group_size=16) - # `frozen` is excluded via _keep_in_fp32_modules; `odd_shape` fails the 128-multiple constraint. - assert targets == ["blocks.0.0", "blocks.1.0", "embedder"] + assert targets == ["blocks.0.proj", "blocks.1.proj"] + + # An explicit list replaces the default name patterns. + targets = infer_data_free_targets(model, group_size=16, modules_to_not_convert=[]) + assert sorted(targets) == ["blocks.0.norm_linear", "blocks.0.proj", "blocks.1.norm_linear", "blocks.1.proj"] - targets = infer_data_free_targets(model, group_size=16, modules_to_not_convert=["embedder"]) - assert targets == ["blocks.0.0", "blocks.1.0"] + targets = infer_data_free_targets(model, group_size=16, modules_to_not_convert=["blocks.0", "norm"]) + assert targets == ["blocks.1.proj"] with pytest.raises(ValueError, match="Could not infer"): - infer_data_free_targets(model, group_size=16, modules_to_not_convert=["blocks", "embedder"]) + infer_data_free_targets(model, group_size=16, modules_to_not_convert=["proj", "norm"]) def test_quantizer_infers_targets_when_omitted(monkeypatch): @@ -309,7 +320,7 @@ def test_quantizer_infers_targets_when_omitted(monkeypatch): config = NunchakuLiteQuantizationConfig( svdq_w4a4={"precision": "nvfp4", "group_size": 16, "rank": 32}, pre_quantized=False, - modules_to_not_convert=["embedder"], + modules_to_not_convert=["norm_linear"], ) quantizer = NunchakuLiteQuantizer(config, pre_quantized=False) model = _InferenceToyModel() @@ -325,7 +336,7 @@ def test_quantizer_infers_targets_when_omitted(monkeypatch): quantizer._process_model_before_weight_loading(model) - assert config.svdq_w4a4["targets"] == ["blocks.0.0", "blocks.1.0"] + assert config.svdq_w4a4["targets"] == ["blocks.0.proj", "blocks.1.proj"] def test_config_targets_optional_only_for_data_free(): From ae79c1b356938d1774ad4aabe7093a8202ada216 Mon Sep 17 00:00:00 2001 From: Pham Hong Vinh Date: Tue, 25 Aug 2026 18:56:25 +0000 Subject: [PATCH 4/7] Rename modules_to_not_convert to exclude_targets Clearer pairing with the svdq_w4a4 `targets` field, and avoids implying the bnb/torchao semantics of keeping modules in high precision at load: the option only filters data-free target inference. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_013NAtDkGmAw2fzbcvjfC79w --- docs/source/en/quantization/nunchaku.md | 2 +- src/diffusers/quantizers/nunchaku/data_free.py | 12 ++++++------ .../quantizers/nunchaku/nunchaku_quantizer.py | 2 +- src/diffusers/quantizers/quantization_config.py | 2 +- tests/quantization/nunchaku/test_data_free.py | 8 ++++---- 5 files changed, 13 insertions(+), 13 deletions(-) diff --git a/docs/source/en/quantization/nunchaku.md b/docs/source/en/quantization/nunchaku.md index e303dba972b9..c149ee300573 100644 --- a/docs/source/en/quantization/nunchaku.md +++ b/docs/source/en/quantization/nunchaku.md @@ -144,7 +144,7 @@ transformer = Flux2Transformer2DModel.from_pretrained( ) ``` -Pass `modules_to_not_convert` (substring match; it replaces the default name patterns) to exclude further modules, or an explicit `targets` list for full control. Quantization happens per weight as the checkpoint streams in, so peak memory stays near the quantized model size. The result is identical to loading a checkpoint produced offline by a data-free SVDQuant exporter. +Pass `exclude_targets` (substring match; it replaces the default name patterns) to exclude further modules, or an explicit `targets` list for full control. Quantization happens per weight as the checkpoint streams in, so peak memory stays near the quantized model size. The result is identical to loading a checkpoint produced offline by a data-free SVDQuant exporter. ## Fused kernels diff --git a/src/diffusers/quantizers/nunchaku/data_free.py b/src/diffusers/quantizers/nunchaku/data_free.py index 4a1f00703b06..a45b9e0b46ed 100644 --- a/src/diffusers/quantizers/nunchaku/data_free.py +++ b/src/diffusers/quantizers/nunchaku/data_free.py @@ -190,7 +190,7 @@ def infer_data_free_targets( model: "torch.nn.Module", *, group_size: int, - modules_to_not_convert: tuple[str, ...] | list[str] | None = None, + exclude_targets: tuple[str, ...] | list[str] | None = None, ) -> list[str]: """Infer quantization targets for data-free mode from a model's structure. @@ -199,15 +199,15 @@ def infer_data_free_targets( divisible by ``group_size``) is selected, restricted to the repeated block stacks of the model (when it has any) so that peripheral modules such as embedders and final projections stay unquantized. Modules whose path - contains a ``modules_to_not_convert`` substring — defaulting to + contains a ``exclude_targets`` substring — defaulting to ``("norm", "modulation")`` to skip adaLN-style linears — or matches the model's ``_keep_in_fp32_modules`` are excluded. Pass an explicit (possibly - empty) ``modules_to_not_convert`` list to replace the default patterns. + empty) ``exclude_targets`` list to replace the default patterns. """ - if modules_to_not_convert is None: - modules_to_not_convert = _DEFAULT_EXCLUDE_PATTERNS - exclude = list(modules_to_not_convert) + list(getattr(model, "_keep_in_fp32_modules", None) or []) + if exclude_targets is None: + exclude_targets = _DEFAULT_EXCLUDE_PATTERNS + exclude = list(exclude_targets) + list(getattr(model, "_keep_in_fp32_modules", None) or []) stack_prefixes = _repeated_block_prefixes(model) targets = [] for name, module in model.named_modules(): diff --git a/src/diffusers/quantizers/nunchaku/nunchaku_quantizer.py b/src/diffusers/quantizers/nunchaku/nunchaku_quantizer.py index 21367573bb4d..9ff8fd0fa5f9 100644 --- a/src/diffusers/quantizers/nunchaku/nunchaku_quantizer.py +++ b/src/diffusers/quantizers/nunchaku/nunchaku_quantizer.py @@ -77,7 +77,7 @@ def _process_model_before_weight_loading( svdq_config["targets"] = infer_data_free_targets( model, group_size=svdq_config["group_size"], - modules_to_not_convert=self.quantization_config.modules_to_not_convert or (), + exclude_targets=self.quantization_config.exclude_targets or (), ) logger.info(f"Inferred {len(svdq_config['targets'])} data-free quantization targets.") diff --git a/src/diffusers/quantizers/quantization_config.py b/src/diffusers/quantizers/quantization_config.py index c41c4f638c16..58b2730aede9 100644 --- a/src/diffusers/quantizers/quantization_config.py +++ b/src/diffusers/quantizers/quantization_config.py @@ -493,7 +493,7 @@ def __init__( raise ValueError("Nunchaku compute_dtype must be a string or a torch.dtype.") self.compute_dtype = compute_dtype self.pre_quantized = kwargs.pop("pre_quantized", True) - self.modules_to_not_convert = kwargs.pop("modules_to_not_convert", None) + self.exclude_targets = kwargs.pop("exclude_targets", None) self.svdq_w4a4 = svdq_w4a4 self.awq_w4a16 = awq_w4a16 diff --git a/tests/quantization/nunchaku/test_data_free.py b/tests/quantization/nunchaku/test_data_free.py index d0b4b52b7997..48a9fc0f3c6b 100644 --- a/tests/quantization/nunchaku/test_data_free.py +++ b/tests/quantization/nunchaku/test_data_free.py @@ -301,14 +301,14 @@ def test_infer_data_free_targets(): assert targets == ["blocks.0.proj", "blocks.1.proj"] # An explicit list replaces the default name patterns. - targets = infer_data_free_targets(model, group_size=16, modules_to_not_convert=[]) + targets = infer_data_free_targets(model, group_size=16, exclude_targets=[]) assert sorted(targets) == ["blocks.0.norm_linear", "blocks.0.proj", "blocks.1.norm_linear", "blocks.1.proj"] - targets = infer_data_free_targets(model, group_size=16, modules_to_not_convert=["blocks.0", "norm"]) + targets = infer_data_free_targets(model, group_size=16, exclude_targets=["blocks.0", "norm"]) assert targets == ["blocks.1.proj"] with pytest.raises(ValueError, match="Could not infer"): - infer_data_free_targets(model, group_size=16, modules_to_not_convert=["proj", "norm"]) + infer_data_free_targets(model, group_size=16, exclude_targets=["proj", "norm"]) def test_quantizer_infers_targets_when_omitted(monkeypatch): @@ -320,7 +320,7 @@ def test_quantizer_infers_targets_when_omitted(monkeypatch): config = NunchakuLiteQuantizationConfig( svdq_w4a4={"precision": "nvfp4", "group_size": 16, "rank": 32}, pre_quantized=False, - modules_to_not_convert=["norm_linear"], + exclude_targets=["norm_linear"], ) quantizer = NunchakuLiteQuantizer(config, pre_quantized=False) model = _InferenceToyModel() From 67da88eba4a6c1ebbb87903eef51b9630ce831a5 Mon Sep 17 00:00:00 2001 From: Pham Hong Vinh Date: Tue, 25 Aug 2026 19:01:48 +0000 Subject: [PATCH 5/7] Silence false missing/unexpected key warnings in data-free mode Match the bitsandbytes loader contract: filter the load-time-produced packed parameter names out of missing_keys via update_missing_keys, and remove the consumed `weight`/`bias` checkpoint keys from unexpected_keys inside create_quantized_param. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_013NAtDkGmAw2fzbcvjfC79w --- .../quantizers/nunchaku/data_free.py | 6 ++++++ .../quantizers/nunchaku/nunchaku_quantizer.py | 11 +++++++++++ tests/quantization/nunchaku/test_data_free.py | 19 +++++++++++++++++++ 3 files changed, 36 insertions(+) diff --git a/src/diffusers/quantizers/nunchaku/data_free.py b/src/diffusers/quantizers/nunchaku/data_free.py index a45b9e0b46ed..5d2c8b297685 100644 --- a/src/diffusers/quantizers/nunchaku/data_free.py +++ b/src/diffusers/quantizers/nunchaku/data_free.py @@ -163,6 +163,12 @@ def pack_lowrank_weight(self, weight: torch.Tensor, down: bool) -> torch.Tensor: return weight.view(c, r) +# Parameter names of SVDQW4A4Linear that data-free quantization produces at +# load time (and therefore never appear in an unquantized checkpoint). +DATA_FREE_PARAMETER_NAMES = frozenset( + {"qweight", "wscales", "wcscales", "wtscale", "smooth_factor", "proj_down", "proj_up"} +) + # adaLN-style linears (feature-wise modulation) are precision-critical and are # consistently named after their norm across diffusers models. _DEFAULT_EXCLUDE_PATTERNS = ("norm", "modulation") diff --git a/src/diffusers/quantizers/nunchaku/nunchaku_quantizer.py b/src/diffusers/quantizers/nunchaku/nunchaku_quantizer.py index 9ff8fd0fa5f9..29ae9cfc61a3 100644 --- a/src/diffusers/quantizers/nunchaku/nunchaku_quantizer.py +++ b/src/diffusers/quantizers/nunchaku/nunchaku_quantizer.py @@ -88,6 +88,15 @@ def _process_model_before_weight_loading( check_strict_state_dict_match(model, state_dict) logger.info(f"Applied Nunchaku quantization config with {num_replaced} targets.") + def update_missing_keys(self, model, missing_keys: list[str], prefix: str) -> list[str]: + if self.pre_quantized: + return missing_keys + # In data-free mode the checkpoint holds `weight`/`bias` while the model + # expects the packed parameters; those are produced at load time. + from .data_free import DATA_FREE_PARAMETER_NAMES + + return [key for key in missing_keys if key.rpartition(".")[2] not in DATA_FREE_PARAMETER_NAMES] + def check_if_quantized_param( self, model: "ModelMixin", @@ -125,6 +134,8 @@ def create_quantized_param( module_name, _, tensor_name = param_name.rpartition(".") module = model.get_submodule(module_name) + if unexpected_keys is not None and param_name in unexpected_keys: + unexpected_keys.remove(param_name) if tensor_name == "bias": packed_bias = pack_data_free_bias(param_value.to(target_device), torch_dtype=self.compute_dtype) module._parameters["bias"] = torch.nn.Parameter(packed_bias, requires_grad=False) diff --git a/tests/quantization/nunchaku/test_data_free.py b/tests/quantization/nunchaku/test_data_free.py index 48a9fc0f3c6b..3be55b794052 100644 --- a/tests/quantization/nunchaku/test_data_free.py +++ b/tests/quantization/nunchaku/test_data_free.py @@ -348,3 +348,22 @@ def test_config_targets_optional_only_for_data_free(): with pytest.raises(ValueError, match="missing required field 'targets'"): NunchakuLiteQuantizationConfig(svdq_w4a4={"precision": "nvfp4", "group_size": 16, "rank": 32}) + + +def test_quantizer_update_missing_keys_filters_data_free_params(): + from diffusers.quantizers.nunchaku.nunchaku_quantizer import NunchakuLiteQuantizer + + config = NunchakuLiteQuantizationConfig( + svdq_w4a4={"precision": "nvfp4", "group_size": 16, "rank": 32}, + pre_quantized=False, + ) + quantizer = NunchakuLiteQuantizer(config, pre_quantized=False) + missing = ["blocks.0.proj.qweight", "blocks.0.proj.smooth_factor", "blocks.0.proj.wtscale", "other.weight"] + assert quantizer.update_missing_keys(None, missing, prefix="") == ["other.weight"] + + pre_config = NunchakuLiteQuantizationConfig( + svdq_w4a4={"precision": "nvfp4", "group_size": 16, "rank": 32, "targets": ["blocks.0.proj"]} + ) + quantizer_pre = NunchakuLiteQuantizer(pre_config, pre_quantized=True) + assert quantizer_pre.pre_quantized is True + assert quantizer_pre.update_missing_keys(None, missing, prefix="") == missing From 6cb4b2e7d3ce36a6d9dd5ede22da41bf06dda4be Mon Sep 17 00:00:00 2001 From: Pham Hong Vinh Date: Wed, 26 Aug 2026 07:20:42 +0000 Subject: [PATCH 6/7] Rename data_free.py to svdquant.py The module implements the SVDQuant math (smoothing, low-rank split, quantization, kernel packing) as opposed to utils.py's kernel runtime; name it after the algorithm. Data-free stays in the function names, where it describes the mode. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_013NAtDkGmAw2fzbcvjfC79w --- src/diffusers/quantizers/nunchaku/nunchaku_quantizer.py | 6 +++--- .../quantizers/nunchaku/{data_free.py => svdquant.py} | 0 .../nunchaku/{test_data_free.py => test_svdquant.py} | 4 ++-- 3 files changed, 5 insertions(+), 5 deletions(-) rename src/diffusers/quantizers/nunchaku/{data_free.py => svdquant.py} (100%) rename tests/quantization/nunchaku/{test_data_free.py => test_svdquant.py} (99%) diff --git a/src/diffusers/quantizers/nunchaku/nunchaku_quantizer.py b/src/diffusers/quantizers/nunchaku/nunchaku_quantizer.py index 29ae9cfc61a3..cf7884b0771e 100644 --- a/src/diffusers/quantizers/nunchaku/nunchaku_quantizer.py +++ b/src/diffusers/quantizers/nunchaku/nunchaku_quantizer.py @@ -72,7 +72,7 @@ def _process_model_before_weight_loading( svdq_config = self.quantization_config.svdq_w4a4 if not self.pre_quantized and svdq_config is not None and svdq_config.get("targets") is None: - from .data_free import infer_data_free_targets + from .svdquant import infer_data_free_targets svdq_config["targets"] = infer_data_free_targets( model, @@ -93,7 +93,7 @@ def update_missing_keys(self, model, missing_keys: list[str], prefix: str) -> li return missing_keys # In data-free mode the checkpoint holds `weight`/`bias` while the model # expects the packed parameters; those are produced at load time. - from .data_free import DATA_FREE_PARAMETER_NAMES + from .svdquant import DATA_FREE_PARAMETER_NAMES return [key for key in missing_keys if key.rpartition(".")[2] not in DATA_FREE_PARAMETER_NAMES] @@ -130,7 +130,7 @@ def create_quantized_param( ): import torch - from .data_free import pack_data_free_bias, quantize_linear_data_free + from .svdquant import pack_data_free_bias, quantize_linear_data_free module_name, _, tensor_name = param_name.rpartition(".") module = model.get_submodule(module_name) diff --git a/src/diffusers/quantizers/nunchaku/data_free.py b/src/diffusers/quantizers/nunchaku/svdquant.py similarity index 100% rename from src/diffusers/quantizers/nunchaku/data_free.py rename to src/diffusers/quantizers/nunchaku/svdquant.py diff --git a/tests/quantization/nunchaku/test_data_free.py b/tests/quantization/nunchaku/test_svdquant.py similarity index 99% rename from tests/quantization/nunchaku/test_data_free.py rename to tests/quantization/nunchaku/test_svdquant.py index 3be55b794052..03a2ee409743 100644 --- a/tests/quantization/nunchaku/test_data_free.py +++ b/tests/quantization/nunchaku/test_svdquant.py @@ -24,7 +24,7 @@ import torch from diffusers import NunchakuLiteQuantizationConfig -from diffusers.quantizers.nunchaku.data_free import ( +from diffusers.quantizers.nunchaku.svdquant import ( _NunchakuWeightPacker, pack_data_free_bias, quantize_linear_data_free, @@ -291,7 +291,7 @@ def __init__(self): def test_infer_data_free_targets(): - from diffusers.quantizers.nunchaku.data_free import infer_data_free_targets + from diffusers.quantizers.nunchaku.svdquant import infer_data_free_targets model = _InferenceToyModel() # Default: restricted to the repeated `blocks` stack (embedder/proj_out are From cae69671ff6502e3b7d0be75467ff880d31f3c8e Mon Sep 17 00:00:00 2001 From: Pham Hong Vinh Date: Wed, 26 Aug 2026 07:37:07 +0000 Subject: [PATCH 7/7] Rename mixin test to test_nunchaku_lite_quantize_on_load Name the integration test after the loader mechanism (pre_quantized=False) rather than the algorithm mode, and rename the companion class attribute to quantize_on_load_config_dict to match. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_013NAtDkGmAw2fzbcvjfC79w --- tests/models/testing_utils/quantization.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/tests/models/testing_utils/quantization.py b/tests/models/testing_utils/quantization.py index 589ef81c5fba..48bf3b283da4 100644 --- a/tests/models/testing_utils/quantization.py +++ b/tests/models/testing_utils/quantization.py @@ -1762,22 +1762,22 @@ def _test_quantized_layers(self, config_kwargs): def test_nunchaku_lite_quantized_layers(self): self._test_quantized_layers(self.config_dict) - def test_nunchaku_lite_data_free_quantization(self): + def test_nunchaku_lite_quantize_on_load(self): """Quantize an unquantized checkpoint on load (`pre_quantized=False`) and run a forward pass.""" unquantized_path = getattr(self, "unquantized_model_name_or_path", None) - data_free_config = getattr(self, "data_free_config_dict", None) - if unquantized_path is None or data_free_config is None: - pytest.skip("Data-free quantization attributes are not configured for this model.") + quantize_on_load_config = getattr(self, "quantize_on_load_config_dict", None) + if unquantized_path is None or quantize_on_load_config is None: + pytest.skip("Quantize-on-load attributes are not configured for this model.") kwargs = getattr(self, "pretrained_model_kwargs", {}).copy() - kwargs["quantization_config"] = NunchakuLiteQuantizationConfig(**data_free_config, pre_quantized=False) + kwargs["quantization_config"] = NunchakuLiteQuantizationConfig(**quantize_on_load_config, pre_quantized=False) model = self.model_class.from_pretrained(unquantized_path, **kwargs) num_quantized_layers = sum(1 for _, module in model.named_modules() if self._is_module_quantized(module)) - expected = len(data_free_config["svdq_w4a4"]["targets"]) + expected = len(quantize_on_load_config["svdq_w4a4"]["targets"]) assert num_quantized_layers == expected, ( - f"Data-free quantization replaced {num_quantized_layers} layers, expected {expected}." + f"Quantize-on-load replaced {num_quantized_layers} layers, expected {expected}." ) with torch.no_grad():