From fc27874a8ea21240bcd21ffd9c23b6ef6fffed0e Mon Sep 17 00:00:00 2001 From: Pham Hong Vinh Date: Tue, 25 Aug 2026 18:40:58 +0000 Subject: [PATCH 01/13] 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 02/13] 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 03/13] 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 04/13] 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 05/13] 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 06/13] 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 07/13] 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(): From c4886e3225fde04d82ba3b417f4457118c7221ce Mon Sep 17 00:00:00 2001 From: Pham Hong Vinh Date: Sat, 29 Aug 2026 01:20:26 +0700 Subject: [PATCH 08/13] Materialize quantize-on-load params for checkpoint keys absent from the model load_model_dict_into_meta skipped any checkpoint key not present in the model's state dict before consulting the quantizer. In data-free mode the replaced SVDQW4A4Linear modules no longer have a `weight` parameter, so the checkpoint's `weight`/`bias` keys were silently dropped, the packed parameters stayed on the meta device, and dispatch_model failed with "Cannot copy out of meta tensor". Give a quantize-on-load quantizer (pre_quantized=False) the chance to claim such keys and materialize the module's packed parameters from them. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01X4NMaGamTgkHjL9EYqNZ7K --- src/diffusers/models/model_loading_utils.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/src/diffusers/models/model_loading_utils.py b/src/diffusers/models/model_loading_utils.py index abbde8082bb5..ec3df5c5a1ab 100644 --- a/src/diffusers/models/model_loading_utils.py +++ b/src/diffusers/models/model_loading_utils.py @@ -234,6 +234,18 @@ def load_model_dict_into_meta( for param_name, param in state_dict.items(): if param_name not in empty_state_dict: + # A quantize-on-load quantizer may claim checkpoint keys that no longer + # exist on the model (e.g. `weight` of a module replaced by a quantized + # linear) and materialize the module's packed parameters from them. + if ( + is_quantized + and not hf_quantizer.pre_quantized + and hf_quantizer.check_if_quantized_param(model, param, param_name, state_dict) + ): + param_device = _determine_param_device(param_name, device_map) + hf_quantizer.create_quantized_param( + model, param, param_name, param_device, state_dict, unexpected_keys, dtype=dtype + ) continue set_module_kwargs = {} From aaa0d21bfeb3d521cc5199d93fadc3b162eff3e5 Mon Sep 17 00:00:00 2001 From: Pham Hong Vinh Date: Sat, 29 Aug 2026 01:46:40 +0700 Subject: [PATCH 09/13] Expose data-free smoothing strength as svdq_w4a4.smooth_exponent The weight-span smoothing exponent was hard-coded to 0.5. Make it an optional config field (default unchanged) so the smoothing strength can be tuned: 0 disables smoothing, 1 fully flattens per-channel weight spans. Only valid in data-free mode. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01X4NMaGamTgkHjL9EYqNZ7K --- .../quantizers/nunchaku/nunchaku_quantizer.py | 1 + src/diffusers/quantizers/nunchaku/svdquant.py | 16 ++++++---- .../quantizers/quantization_config.py | 11 +++++++ tests/quantization/nunchaku/test_svdquant.py | 29 +++++++++++++++++++ 4 files changed, 52 insertions(+), 5 deletions(-) diff --git a/src/diffusers/quantizers/nunchaku/nunchaku_quantizer.py b/src/diffusers/quantizers/nunchaku/nunchaku_quantizer.py index cf7884b0771e..6741b2801e96 100644 --- a/src/diffusers/quantizers/nunchaku/nunchaku_quantizer.py +++ b/src/diffusers/quantizers/nunchaku/nunchaku_quantizer.py @@ -146,6 +146,7 @@ def create_quantized_param( group_size=module.group_size, rank=module.rank, torch_dtype=self.compute_dtype, + smooth_exponent=self.quantization_config.svdq_w4a4.get("smooth_exponent", 0.5), ) for name, tensor in quantized.items(): module._parameters[name] = torch.nn.Parameter(tensor.to(target_device), requires_grad=False) diff --git a/src/diffusers/quantizers/nunchaku/svdquant.py b/src/diffusers/quantizers/nunchaku/svdquant.py index 5d2c8b297685..0b31ea917ef8 100644 --- a/src/diffusers/quantizers/nunchaku/svdquant.py +++ b/src/diffusers/quantizers/nunchaku/svdquant.py @@ -246,15 +246,19 @@ def _check_packable(out_features: int, in_features: int, rank: int, group_size: 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``. +def _weight_span_smooth_scale(weight: torch.Tensor, exponent: float = 0.5) -> torch.Tensor: + """Data-free weight-span smoothing: ``s_j = 1 / absmax(W[:, j]) ** exponent``. The weight is stored multiplied by ``s`` (equalizing per-channel magnitudes) - and the kernel divides the activations by ``s`` at runtime. + and the kernel divides the activations by ``s`` at runtime. ``exponent`` + controls how strongly channel magnitudes are equalized: ``0`` disables + smoothing, ``1`` fully flattens the per-channel weight spans. Without + calibration data there is no activation term (SmoothQuant's alpha), so this + weight-side exponent is the only smoothing knob. """ span = weight.abs().amax(dim=0).clamp_min(_SMOOTH_EPS) - scale = 1.0 / span.pow(0.5) + scale = 1.0 / span.pow(exponent) scale = torch.where(torch.isfinite(scale), scale, torch.ones_like(scale)) return scale.clamp_min(_SMOOTH_EPS) @@ -273,6 +277,7 @@ def quantize_linear_data_free( group_size: int, rank: int, torch_dtype: torch.dtype = torch.bfloat16, + smooth_exponent: float = 0.5, ) -> dict[str, torch.Tensor]: """Quantize one linear weight into ``SVDQW4A4Linear``'s packed parameters. @@ -282,6 +287,7 @@ def quantize_linear_data_free( 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. + smooth_exponent: Weight-span smoothing strength in ``[0, 1]``. Returns: Mapping with keys ``qweight``, ``wscales``, ``smooth_factor``, @@ -293,7 +299,7 @@ def quantize_linear_data_free( packer = _NunchakuWeightPacker() weight = weight.to(dtype=torch.float32) - smooth = _weight_span_smooth_scale(weight) + smooth = _weight_span_smooth_scale(weight, exponent=smooth_exponent) smoothed = weight * smooth.view(1, -1) if rank > 0: diff --git a/src/diffusers/quantizers/quantization_config.py b/src/diffusers/quantizers/quantization_config.py index 58b2730aede9..e167243c02df 100644 --- a/src/diffusers/quantizers/quantization_config.py +++ b/src/diffusers/quantizers/quantization_config.py @@ -555,6 +555,17 @@ def post_init(self): f"Nunchaku SVDQ config with precision={precision!r} requires " f"group_size={expected_group_size}, got {group_size}." ) + if "smooth_exponent" in raw: + smooth_exponent = raw["smooth_exponent"] + if self.pre_quantized: + raise ValueError( + "'smooth_exponent' only applies to data-free quantization (`pre_quantized=False`); " + "a pre-quantized checkpoint already has its smoothing baked in." + ) + if not isinstance(smooth_exponent, (int, float)) or isinstance(smooth_exponent, bool): + raise ValueError(f"Nunchaku compact config section {op!r} field 'smooth_exponent' must be a number.") + if not 0.0 <= smooth_exponent <= 1.0: + raise ValueError(f"'smooth_exponent' must be in [0, 1], got {smooth_exponent}.") elif precision != "int4": raise ValueError("Nunchaku AWQ target requires precision='int4'.") diff --git a/tests/quantization/nunchaku/test_svdquant.py b/tests/quantization/nunchaku/test_svdquant.py index 03a2ee409743..366808a7bca1 100644 --- a/tests/quantization/nunchaku/test_svdquant.py +++ b/tests/quantization/nunchaku/test_svdquant.py @@ -339,6 +339,35 @@ def test_quantizer_infers_targets_when_omitted(monkeypatch): assert config.svdq_w4a4["targets"] == ["blocks.0.proj", "blocks.1.proj"] +def test_smooth_exponent_flows_through(): + torch.manual_seed(0) + weight = torch.randn(128, 128) + kwargs = dict(precision="int4", group_size=64, rank=0) + default = quantize_linear_data_free(weight, **kwargs) + strong = quantize_linear_data_free(weight, smooth_exponent=1.0, **kwargs) + disabled = quantize_linear_data_free(weight, smooth_exponent=0.0, **kwargs) + # Exponent 0 disables smoothing entirely; other exponents change the factor. + assert torch.equal(disabled["smooth_factor"], torch.ones_like(disabled["smooth_factor"])) + assert not torch.equal(strong["smooth_factor"], default["smooth_factor"]) + + +def test_config_validates_smooth_exponent(): + section = {"precision": "nvfp4", "group_size": 16, "rank": 32, "smooth_exponent": 0.25} + config = NunchakuLiteQuantizationConfig(svdq_w4a4=dict(section), pre_quantized=False) + assert config.svdq_w4a4["smooth_exponent"] == 0.25 + + with pytest.raises(ValueError, match="must be in \\[0, 1\\]"): + NunchakuLiteQuantizationConfig( + svdq_w4a4={**section, "smooth_exponent": 1.5}, pre_quantized=False + ) + with pytest.raises(ValueError, match="must be a number"): + NunchakuLiteQuantizationConfig( + svdq_w4a4={**section, "smooth_exponent": "0.5"}, pre_quantized=False + ) + with pytest.raises(ValueError, match="only applies to data-free"): + NunchakuLiteQuantizationConfig(svdq_w4a4={**section, "targets": ["proj"]}) + + def test_config_targets_optional_only_for_data_free(): config = NunchakuLiteQuantizationConfig( svdq_w4a4={"precision": "nvfp4", "group_size": 16, "rank": 32}, From aa4428da1f69416d5bff46c16a33b9a09464702c Mon Sep 17 00:00:00 2001 From: Pham Hong Vinh Date: Mon, 31 Aug 2026 13:04:26 +0700 Subject: [PATCH 10/13] Document rank and smooth_exponent tuning for data-free quantization Measured on a 19B video DiT: int4 at rank=128 recovers ~2.7 dB over rank=32 for ~16% more transformer memory, and the optimal smoothing strength decreases as rank grows - note the interaction so users sweep the exponent when raising rank instead of assuming the 0.5 default. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01X4NMaGamTgkHjL9EYqNZ7K --- docs/source/en/quantization/nunchaku.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/docs/source/en/quantization/nunchaku.md b/docs/source/en/quantization/nunchaku.md index c149ee300573..554cac54e31e 100644 --- a/docs/source/en/quantization/nunchaku.md +++ b/docs/source/en/quantization/nunchaku.md @@ -146,6 +146,10 @@ transformer = Flux2Transformer2DModel.from_pretrained( 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. +Two quality knobs are worth tuning per model. `rank` sets the size of the bf16 low-rank branch that absorbs the largest singular components before 4-bit coding — larger ranks trade a modest size increase for fidelity, and the gain can be substantial (on a 19B video DiT, int4 at `rank=128` recovered ~2.7 dB over `rank=32` for ~16% more transformer memory). `smooth_exponent` (default `0.5`) sets the weight-span smoothing strength; it interacts with `rank` — larger low-rank branches tend to prefer weaker smoothing (e.g. `0.25` at `rank=128`) — so when raising `rank`, sweep the exponent rather than assuming the default. + +Quantize-on-load runs one SVD per target at load time (minutes for large models on a fast GPU). For repeated loads of the same configuration, prefer a pre-quantized checkpoint produced by an offline exporter — loading packed weights takes seconds, and the packed format is identical. + ## 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. From 93f6cc22d6467accd8ecf78609e200e778546d4f Mon Sep 17 00:00:00 2001 From: Pham Hong Vinh Date: Tue, 1 Sep 2026 20:09:36 +0700 Subject: [PATCH 11/13] Allow CPU quantize-on-load and save_pretrained for Nunchaku Lite Data-free SVDQuant (weight-span smoothing, SVD, group quantization) is pure PyTorch and never invokes a CUDA kernel, but validate_environment unconditionally required a GPU and the kernels package, and quantizers/nunchaku/utils.py fetched the CUDA kernel package eagerly at import time even though `ops` is only used by SVDQW4A4Linear.forward at real inference time. Both are now scoped to when they actually matter: - validate_environment returns early for `pre_quantized=False` loads with no CUDA available; the real-inference (`pre_quantized=True`) path is unchanged. - `ops` in nunchaku/utils.py is now a lazy proxy that defers the kernels-package/network fetch to the first actual kernel call, so the module imports without the `kernels` package or a GPU. - is_serializable is now True. `_process_model_before_weight_loading` builds a save-ready config after target inference (pre_quantized=True, data-free-only knobs like smooth_exponent stripped - post_init rejects them once pre_quantized is True) and assigns it as both a dict item and an attribute on `model.config`, since `model.config` is a FrozenDict whose attribute assignment only shadows via instance __dict__ and is invisible to save_pretrained's dict-based serialization. Together this lets a GPU-less machine quantize-on-load and `save_pretrained` a checkpoint that loads normally via the existing pre_quantized=True path elsewhere. Verified with a real (non-stubbed) round trip: CPU quantize -> save_pretrained -> reload, packed tensors bit-identical. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01X4NMaGamTgkHjL9EYqNZ7K --- docs/source/en/quantization/nunchaku.md | 26 ++++++ .../quantizers/nunchaku/nunchaku_quantizer.py | 41 +++++++++- src/diffusers/quantizers/nunchaku/utils.py | 36 +++++++-- tests/quantization/nunchaku/test_svdquant.py | 79 +++++++++++++++++++ 4 files changed, 173 insertions(+), 9 deletions(-) diff --git a/docs/source/en/quantization/nunchaku.md b/docs/source/en/quantization/nunchaku.md index 554cac54e31e..f49c8001d7ad 100644 --- a/docs/source/en/quantization/nunchaku.md +++ b/docs/source/en/quantization/nunchaku.md @@ -150,6 +150,32 @@ Two quality knobs are worth tuning per model. `rank` sets the size of the bf16 l Quantize-on-load runs one SVD per target at load time (minutes for large models on a fast GPU). For repeated loads of the same configuration, prefer a pre-quantized checkpoint produced by an offline exporter — loading packed weights takes seconds, and the packed format is identical. +### Quantize on a CPU-only machine, then save + +Data-free quantization is pure PyTorch (weight-span smoothing, SVD, group quantization) and never calls a CUDA kernel, so it also runs on a machine with no GPU and no `kernels` package installed. Quantize once, `save_pretrained` the result, then load the packed checkpoint anywhere with the usual (fast, kernel-free) `pre_quantized` path: + +```python +import torch +from diffusers import Flux2Transformer2DModel, NunchakuLiteQuantizationConfig + +# Runs on CPU - no GPU or `kernels` package required for this step. +transformer = Flux2Transformer2DModel.from_pretrained( + "black-forest-labs/FLUX.2-klein-9B", + subfolder="transformer", + quantization_config=NunchakuLiteQuantizationConfig( + svdq_w4a4={"precision": "nvfp4", "group_size": 16, "rank": 32}, + pre_quantized=False, + ), + torch_dtype=torch.bfloat16, +) +transformer.save_pretrained("flux2-klein-9b-nvfp4") + +# Later, on a GPU machine, loads through the ordinary pre-quantized path. +transformer = Flux2Transformer2DModel.from_pretrained( + "flux2-klein-9b-nvfp4", torch_dtype=torch.bfloat16, device_map="cuda" +) +``` + ## 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/nunchaku_quantizer.py b/src/diffusers/quantizers/nunchaku/nunchaku_quantizer.py index 6741b2801e96..443adc7b5daa 100644 --- a/src/diffusers/quantizers/nunchaku/nunchaku_quantizer.py +++ b/src/diffusers/quantizers/nunchaku/nunchaku_quantizer.py @@ -26,12 +26,20 @@ def __init__(self, quantization_config, **kwargs): self.pre_quantized = self.pre_quantized and quantization_config.pre_quantized def validate_environment(self, *args, **kwargs): + import torch + + if not self.pre_quantized and not torch.cuda.is_available(): + # Data-free quantize-on-load (weight-span smoothing, SVD, group + # quantization) is pure PyTorch and never invokes a CUDA kernel, so a + # GPU-less machine can quantize-and-export a checkpoint here for later + # GPU inference (see `save_pretrained`). + return + if not is_kernels_available(): raise ImportError( "Loading Nunchaku checkpoints requires the Hugging Face `kernels` package. " "Install it with `pip install kernels`." ) - import torch cuda_available = torch.cuda.is_available() if not cuda_available: @@ -86,6 +94,35 @@ def _process_model_before_weight_loading( if self.pre_quantized and state_dict is not None: check_strict_state_dict_match(model, state_dict) + + if not self.pre_quantized: + # `save_pretrained` afterward must write a config a later `pre_quantized=True` + # reload accepts: `smooth_exponent` and other data-free-only knobs are baked + # into the packed residual and are rejected by `post_init` once pre_quantized + # is True, and the checkpoint itself is now pre-quantized either way. Building + # a fresh config (rather than mutating `self.quantization_config`, which this + # load's own `pre_quantized=False` state must keep) also re-validates it. + from ..quantization_config import NunchakuLiteQuantizationConfig + + export_config = NunchakuLiteQuantizationConfig( + compute_dtype=self.compute_dtype, + svdq_w4a4={ + "precision": svdq_config["precision"], + "group_size": svdq_config["group_size"], + "rank": svdq_config["rank"], + "targets": svdq_config["targets"], + }, + pre_quantized=True, + ) + # `model.config` is a `FrozenDict`: attribute assignment only shadows via + # instance `__dict__` and is invisible to the dict-based serialization + # `save_pretrained` uses, so the item assignment below is the one that + # actually matters; the attribute is set too so in-memory reads of + # `model.config.quantization_config` (e.g. right after this load) agree. + if hasattr(model.config, "__setitem__"): + model.config["quantization_config"] = export_config + model.config.quantization_config = export_config + 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]: @@ -156,7 +193,7 @@ def _process_model_after_weight_loading(self, model: "ModelMixin", **kwargs): @property def is_serializable(self): - return False + return True @property def is_trainable(self) -> bool: diff --git a/src/diffusers/quantizers/nunchaku/utils.py b/src/diffusers/quantizers/nunchaku/utils.py index 7f220dfc7e0f..c03ba1adbd58 100644 --- a/src/diffusers/quantizers/nunchaku/utils.py +++ b/src/diffusers/quantizers/nunchaku/utils.py @@ -21,7 +21,12 @@ _HF_KERNEL_VERSION = 2 -if is_kernels_available(): +def _fetch_kernel_ops(): + if not is_kernels_available(): + raise ImportError( + "Loading Nunchaku checkpoints requires the Hugging Face `kernels` package. " + "Install it with `pip install kernels`." + ) from kernels import get_kernel if not DIFFUSERS_TRUST_REMOTE_KERNELS: @@ -31,14 +36,31 @@ ) # `kernels<0.14.0` has no `trust_remote_code` argument and executes the downloaded code unconditionally. trust_kwargs = {"trust_remote_code": True} if is_kernels_version(">=", "0.14.0") else {} - ops = get_kernel( + return get_kernel( _HF_KERNEL_REPO, version=_HF_KERNEL_VERSION, user_agent={"diffusers": __version__}, **trust_kwargs ).ops -else: - raise ImportError( - "Loading Nunchaku checkpoints requires the Hugging Face `kernels` package. " - "Install it with `pip install kernels`." - ) + + +class _LazyKernelOps: + """Defers the kernels package/network fetch to the first actual kernel call. + + Constructing `SVDQW4A4Linear`/`AWQW4A16Linear` modules and running data-free + quantize-on-load (`quantize_linear_data_free`, pure PyTorch) never touch these + ops; only a packed module's `.forward()` does. This lets `nunchaku/utils.py` be + imported - and a checkpoint quantized and exported - on a machine with no GPU + and no `kernels` package installed, while keeping identical error behavior + (`ImportError`/`ValueError`) once inference is actually attempted. + """ + + _ops = None + + def __getattr__(self, name): + if self._ops is None: + self._ops = _fetch_kernel_ops() + return getattr(self._ops, name) + + +ops = _LazyKernelOps() def _gemm_w4a4( diff --git a/tests/quantization/nunchaku/test_svdquant.py b/tests/quantization/nunchaku/test_svdquant.py index 366808a7bca1..1cdff32c1838 100644 --- a/tests/quantization/nunchaku/test_svdquant.py +++ b/tests/quantization/nunchaku/test_svdquant.py @@ -324,6 +324,7 @@ def test_quantizer_infers_targets_when_omitted(monkeypatch): ) quantizer = NunchakuLiteQuantizer(config, pre_quantized=False) model = _InferenceToyModel() + model.config = types.SimpleNamespace() # plain nn.Module here; ModelMixin normally provides this # Stub out `.utils` (its import fetches the CUDA kernels) so only the # target-inference part of _process_model_before_weight_loading runs. @@ -396,3 +397,81 @@ def test_quantizer_update_missing_keys_filters_data_free_params(): 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 + + +def test_validate_environment_allows_cpu_for_data_free(monkeypatch): + import torch + + import diffusers.quantizers.nunchaku.nunchaku_quantizer as nunchaku_quantizer_module + from diffusers.quantizers.nunchaku.nunchaku_quantizer import NunchakuLiteQuantizer + + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + # `kernels` isn't installed in this CPU-only test environment either; that must + # not matter for the data-free bypass (it returns before that check), but pin + # it explicitly so the pre_quantized=True branch below hits the intended + # CUDA-capable error rather than an environment-dependent ImportError. + monkeypatch.setattr(nunchaku_quantizer_module, "is_kernels_available", lambda: True) + + data_free_config = NunchakuLiteQuantizationConfig( + svdq_w4a4={"precision": "nvfp4", "group_size": 16, "rank": 32}, + pre_quantized=False, + ) + quantizer = NunchakuLiteQuantizer(data_free_config, pre_quantized=False) + quantizer.validate_environment() # must not raise, even without `kernels` installed + + pre_quantized_config = NunchakuLiteQuantizationConfig( + svdq_w4a4={"precision": "nvfp4", "group_size": 16, "rank": 32, "targets": ["proj"]} + ) + quantizer_pre = NunchakuLiteQuantizer(pre_quantized_config, pre_quantized=True) + with pytest.raises(ValueError, match="CUDA-capable"): + quantizer_pre.validate_environment() + + +def test_is_serializable(): + from diffusers.quantizers.nunchaku.nunchaku_quantizer import NunchakuLiteQuantizer + + config = NunchakuLiteQuantizationConfig( + svdq_w4a4={"precision": "nvfp4", "group_size": 16, "rank": 32, "targets": ["proj"]} + ) + assert NunchakuLiteQuantizer(config, pre_quantized=True).is_serializable is True + + +def test_quantizer_exports_reload_ready_config_after_data_free_load(monkeypatch): + import sys + import types + + from diffusers.quantizers.nunchaku.nunchaku_quantizer import NunchakuLiteQuantizer + + config = NunchakuLiteQuantizationConfig( + svdq_w4a4={"precision": "nvfp4", "group_size": 16, "rank": 32, "smooth_exponent": 0.25}, + pre_quantized=False, + # Explicit (matching the sibling `test_quantizer_infers_targets_when_omitted` + # test): `exclude_targets=None` on the config becomes `()` in the call to + # `infer_data_free_targets`, which is itself the "no exclusions" override + # rather than "use the norm/modulation defaults" - a separate, pre-existing + # quirk this test sidesteps rather than depends on. + exclude_targets=["norm_linear"], + ) + quantizer = NunchakuLiteQuantizer(config, pre_quantized=False) + model = _InferenceToyModel() + model.config = types.SimpleNamespace() # plain nn.Module here; ModelMixin normally provides this + + # Stub out `.utils` (its import fetches the CUDA kernels) so only the + # target-inference and config-export parts of _process_model_before_weight_loading run. + 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) + + exported = model.config.quantization_config + assert isinstance(exported, NunchakuLiteQuantizationConfig) + assert exported.pre_quantized is True + assert exported.svdq_w4a4["targets"] == ["blocks.0.proj", "blocks.1.proj"] + assert "smooth_exponent" not in exported.svdq_w4a4 + # The original, still-loading config keeps its own pre_quantized=False state. + assert quantizer.pre_quantized is False + assert config.pre_quantized is False From 8ffd5ef2ec2519a8e5eefec0a31c564a1be54797 Mon Sep 17 00:00:00 2001 From: Pham Hong Vinh Date: Tue, 1 Sep 2026 20:17:18 +0700 Subject: [PATCH 12/13] Stop coercing exclude_targets=None into () before target inference _process_model_before_weight_loading passed `self.quantization_config.exclude_targets or ()` to infer_data_free_targets, turning the config's own None default into an explicit empty sequence. infer_data_free_targets treats None as "use the norm/modulation defaults" and any explicit sequence (including an empty one) as "use exactly this instead" - so a config that never sets exclude_targets silently got zero exclusions through the quantizer, unlike calling infer_data_free_targets directly with the argument genuinely omitted. Pass the config's value through as-is and let its own None-handling apply. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01X4NMaGamTgkHjL9EYqNZ7K --- .../quantizers/nunchaku/nunchaku_quantizer.py | 6 ++- tests/quantization/nunchaku/test_svdquant.py | 49 ++++++++++++++++--- 2 files changed, 48 insertions(+), 7 deletions(-) diff --git a/src/diffusers/quantizers/nunchaku/nunchaku_quantizer.py b/src/diffusers/quantizers/nunchaku/nunchaku_quantizer.py index 443adc7b5daa..e25058f4bd4d 100644 --- a/src/diffusers/quantizers/nunchaku/nunchaku_quantizer.py +++ b/src/diffusers/quantizers/nunchaku/nunchaku_quantizer.py @@ -85,7 +85,11 @@ def _process_model_before_weight_loading( svdq_config["targets"] = infer_data_free_targets( model, group_size=svdq_config["group_size"], - exclude_targets=self.quantization_config.exclude_targets or (), + # `None` (the config's own default when unset) must reach + # `infer_data_free_targets` as-is, not coerced to `()` - that function + # treats `None` as "use the norm/modulation defaults" and any explicit + # sequence, including an empty one, as "use exactly this instead". + exclude_targets=self.quantization_config.exclude_targets, ) logger.info(f"Inferred {len(svdq_config['targets'])} data-free quantization targets.") diff --git a/tests/quantization/nunchaku/test_svdquant.py b/tests/quantization/nunchaku/test_svdquant.py index 1cdff32c1838..570427417061 100644 --- a/tests/quantization/nunchaku/test_svdquant.py +++ b/tests/quantization/nunchaku/test_svdquant.py @@ -445,12 +445,9 @@ def test_quantizer_exports_reload_ready_config_after_data_free_load(monkeypatch) config = NunchakuLiteQuantizationConfig( svdq_w4a4={"precision": "nvfp4", "group_size": 16, "rank": 32, "smooth_exponent": 0.25}, pre_quantized=False, - # Explicit (matching the sibling `test_quantizer_infers_targets_when_omitted` - # test): `exclude_targets=None` on the config becomes `()` in the call to - # `infer_data_free_targets`, which is itself the "no exclusions" override - # rather than "use the norm/modulation defaults" - a separate, pre-existing - # quirk this test sidesteps rather than depends on. - exclude_targets=["norm_linear"], + # exclude_targets omitted: relies on the default norm/modulation exclusion + # patterns to filter out `norm_linear` (see + # test_quantizer_applies_default_exclude_patterns_when_omitted). ) quantizer = NunchakuLiteQuantizer(config, pre_quantized=False) model = _InferenceToyModel() @@ -475,3 +472,43 @@ def test_quantizer_exports_reload_ready_config_after_data_free_load(monkeypatch) # The original, still-loading config keeps its own pre_quantized=False state. assert quantizer.pre_quantized is False assert config.pre_quantized is False + + +def test_quantizer_applies_default_exclude_patterns_when_omitted(monkeypatch): + """A config with no `exclude_targets` must still get the norm/modulation defaults. + + Regression test: `_process_model_before_weight_loading` used to pass + `self.quantization_config.exclude_targets or ()` to `infer_data_free_targets`, + turning the config's own `None` default into an explicit empty sequence - + which that function treats as "use exactly this (no exclusions)" rather than + "fall back to the norm/modulation defaults". `norm_linear`-style targets were + silently never excluded when going through the quantizer, unlike a direct + `infer_data_free_targets(model, group_size=...)` call with the argument + genuinely omitted (see `test_infer_data_free_targets`). + """ + 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, + ) + assert config.exclude_targets is None + quantizer = NunchakuLiteQuantizer(config, pre_quantized=False) + model = _InferenceToyModel() + model.config = types.SimpleNamespace() + + 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) + + # Matches infer_data_free_targets(model, group_size=16)'s own default-pattern + # result in test_infer_data_free_targets: norm_linear is excluded. + assert config.svdq_w4a4["targets"] == ["blocks.0.proj", "blocks.1.proj"] From f0ed9073b3f4c20be15959bfa58edc3630a672c8 Mon Sep 17 00:00:00 2001 From: Pham Hong Vinh Date: Tue, 1 Sep 2026 20:29:08 +0700 Subject: [PATCH 13/13] Show pipeline-level save/load in the CPU-quantize doc example transformer.save_pretrained() only writes the bare transformer, which isn't loadable via DiffusionPipeline.from_pretrained the way "Load a quantized pipeline" above it demonstrates. DiffusionPipeline.save_pretrained just calls each registered component's own save_pretrained (verified: the quantized transformer round-trips bit-identical through it, same as the model-level path), so save the whole pipeline instead - one call produces a self-contained, directly reloadable checkpoint. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01X4NMaGamTgkHjL9EYqNZ7K --- docs/source/en/quantization/nunchaku.md | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/docs/source/en/quantization/nunchaku.md b/docs/source/en/quantization/nunchaku.md index f49c8001d7ad..7dabd1757e97 100644 --- a/docs/source/en/quantization/nunchaku.md +++ b/docs/source/en/quantization/nunchaku.md @@ -152,15 +152,17 @@ Quantize-on-load runs one SVD per target at load time (minutes for large models ### Quantize on a CPU-only machine, then save -Data-free quantization is pure PyTorch (weight-span smoothing, SVD, group quantization) and never calls a CUDA kernel, so it also runs on a machine with no GPU and no `kernels` package installed. Quantize once, `save_pretrained` the result, then load the packed checkpoint anywhere with the usual (fast, kernel-free) `pre_quantized` path: +Data-free quantization is pure PyTorch (weight-span smoothing, SVD, group quantization) and never calls a CUDA kernel, so it also runs on a machine with no GPU and no `kernels` package installed. Quantize the transformer once, `save_pretrained` the whole pipeline, and load the result anywhere with the usual (fast, kernel-free) `pre_quantized` path — [`~DiffusionPipeline.save_pretrained`] saves every component to its own subfolder (each by calling that component's own `save_pretrained`), so this produces the same self-contained, directly loadable layout as [Load a quantized pipeline](#load-a-quantized-pipeline) above: ```python import torch -from diffusers import Flux2Transformer2DModel, NunchakuLiteQuantizationConfig +from diffusers import Flux2Pipeline, Flux2Transformer2DModel, NunchakuLiteQuantizationConfig + +model_id = "black-forest-labs/FLUX.2-klein-9B" # Runs on CPU - no GPU or `kernels` package required for this step. transformer = Flux2Transformer2DModel.from_pretrained( - "black-forest-labs/FLUX.2-klein-9B", + model_id, subfolder="transformer", quantization_config=NunchakuLiteQuantizationConfig( svdq_w4a4={"precision": "nvfp4", "group_size": 16, "rank": 32}, @@ -168,12 +170,11 @@ transformer = Flux2Transformer2DModel.from_pretrained( ), torch_dtype=torch.bfloat16, ) -transformer.save_pretrained("flux2-klein-9b-nvfp4") +pipe = Flux2Pipeline.from_pretrained(model_id, transformer=transformer, torch_dtype=torch.bfloat16) +pipe.save_pretrained("flux2-klein-9b-nvfp4") # Later, on a GPU machine, loads through the ordinary pre-quantized path. -transformer = Flux2Transformer2DModel.from_pretrained( - "flux2-klein-9b-nvfp4", torch_dtype=torch.bfloat16, device_map="cuda" -) +pipe = Flux2Pipeline.from_pretrained("flux2-klein-9b-nvfp4", torch_dtype=torch.bfloat16).to("cuda") ``` ## Fused kernels