diff --git a/docs/source/en/quantization/nunchaku.md b/docs/source/en/quantization/nunchaku.md index d6a07591b6fe..c149ee300573 100644 --- a/docs/source/en/quantization/nunchaku.md +++ b/docs/source/en/quantization/nunchaku.md @@ -122,6 +122,30 @@ 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). + +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 +from diffusers import Flux2Transformer2DModel, NunchakuLiteQuantizationConfig + +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, + device_map="cuda", +) +``` + +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 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 b8f20d7ddba3..cf7884b0771e 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(): @@ -66,13 +70,86 @@ 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 .svdquant import infer_data_free_targets + + svdq_config["targets"] = infer_data_free_targets( + model, + group_size=svdq_config["group_size"], + exclude_targets=self.quantization_config.exclude_targets 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) - 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 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 .svdquant 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", + 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 .svdquant import pack_data_free_bias, quantize_linear_data_free + + 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) + 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/nunchaku/svdquant.py b/src/diffusers/quantizers/nunchaku/svdquant.py new file mode 100644 index 000000000000..5d2c8b297685 --- /dev/null +++ b/src/diffusers/quantizers/nunchaku/svdquant.py @@ -0,0 +1,345 @@ +"""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) + + +# 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") + + +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, + exclude_targets: 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, 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 ``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) ``exclude_targets`` list to replace the default patterns. + """ + + 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(): + 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: + 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( + "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/quantization_config.py b/src/diffusers/quantizers/quantization_config.py index ea78b5f7ff53..58b2730aede9 100644 --- a/src/diffusers/quantizers/quantization_config.py +++ b/src/diffusers/quantizers/quantization_config.py @@ -492,7 +492,8 @@ 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.exclude_targets = kwargs.pop("exclude_targets", None) self.svdq_w4a4 = svdq_w4a4 self.awq_w4a16 = awq_w4a16 @@ -503,6 +504,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: @@ -510,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( @@ -520,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/models/testing_utils/quantization.py b/tests/models/testing_utils/quantization.py index 918126fe3f13..48bf3b283da4 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_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) + 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(**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(quantize_on_load_config["svdq_w4a4"]["targets"]) + assert num_quantized_layers == expected, ( + f"Quantize-on-load 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_svdquant.py b/tests/quantization/nunchaku/test_svdquant.py new file mode 100644 index 000000000000..03a2ee409743 --- /dev/null +++ b/tests/quantization/nunchaku/test_svdquant.py @@ -0,0 +1,369 @@ +# 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.svdquant 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,) + + +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([_ToyBlock() for _ in range(2)]) + self.embedder = torch.nn.Linear(IN_FEATURES, OUT_FEATURES) + self.proj_out = torch.nn.Linear(OUT_FEATURES, IN_FEATURES) + + +def test_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 + # outside), minus adaLN-style names ("norm"), _keep_in_fp32_modules, and + # dimension-ineligible layers. + targets = infer_data_free_targets(model, group_size=16) + 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, 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, 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, exclude_targets=["proj", "norm"]) + + +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, + exclude_targets=["norm_linear"], + ) + 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.proj", "blocks.1.proj"] + + +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}) + + +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