From 82f60db187075e7ef4c331aac8e94a38a028ac34 Mon Sep 17 00:00:00 2001 From: changwangss Date: Tue, 28 Jul 2026 16:14:01 +0800 Subject: [PATCH 1/9] feat: port SVDQuant core primitives Signed-off-by: changwangss --- .../transforms/svdquant/__init__.py | 18 ++ .../algorithms/transforms/svdquant/config.py | 97 ++++++++ .../transforms/svdquant/residual.py | 231 ++++++++++++++++++ .../algorithms/transforms/svdquant/wrapper.py | 42 ++++ .../algorithms/test_svdquant_residual.py | 152 ++++++++++++ 5 files changed, 540 insertions(+) create mode 100644 auto_round/algorithms/transforms/svdquant/__init__.py create mode 100644 auto_round/algorithms/transforms/svdquant/config.py create mode 100644 auto_round/algorithms/transforms/svdquant/residual.py create mode 100644 auto_round/algorithms/transforms/svdquant/wrapper.py create mode 100644 test/test_cpu/algorithms/test_svdquant_residual.py diff --git a/auto_round/algorithms/transforms/svdquant/__init__.py b/auto_round/algorithms/transforms/svdquant/__init__.py new file mode 100644 index 0000000000..14099b8a10 --- /dev/null +++ b/auto_round/algorithms/transforms/svdquant/__init__.py @@ -0,0 +1,18 @@ +# Copyright (c) 2026 Intel Corporation +# +# 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. + +from auto_round.algorithms.transforms.svdquant.config import SVDQuantConfig +from auto_round.algorithms.transforms.svdquant.wrapper import SVDQuantLinear + +__all__ = ["SVDQuantConfig", "SVDQuantLinear"] diff --git a/auto_round/algorithms/transforms/svdquant/config.py b/auto_round/algorithms/transforms/svdquant/config.py new file mode 100644 index 0000000000..16719a6e40 --- /dev/null +++ b/auto_round/algorithms/transforms/svdquant/config.py @@ -0,0 +1,97 @@ +# Copyright (c) 2026 Intel Corporation +# +# 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. + +from __future__ import annotations + +from auto_round.algorithms.quantization.config import QuantizationConfig + + +class SVDQuantConfig(QuantizationConfig): + """Configuration for the SVDQuant structural preprocessor.""" + + def __init__( + self, + *, + rank: int = 32, + smooth_enabled: bool = False, + smooth_num_grids: int = 20, + smooth_max_calibration_calls: int = 128, + target_modules: list[str] | tuple[str, ...] | str | None = None, + exclude_modules: list[str] | tuple[str, ...] | str | None = None, + low_rank_dtype: str = "bf16", + smooth_eps: float = 1e-6, + residual_iters: int = 1, + residual_early_stop: bool = False, + residual_quant_method: str = "rtn", + model_adapter: str | None = None, + **kwargs, + ): + super().__init__(**kwargs) + residual_quant_method = residual_quant_method.lower() + if type(smooth_enabled) is not bool: + raise ValueError(f"`smooth_enabled` must be a bool, got {smooth_enabled!r}") + if not isinstance(rank, int) or isinstance(rank, bool) or rank < 0: + raise ValueError(f"`rank` must be a non-negative integer, got {rank!r}") + if type(smooth_num_grids) is not int or smooth_num_grids < 2: + raise ValueError( + f"`smooth_num_grids` must be an integer greater than or equal to 2, got {smooth_num_grids!r}" + ) + if type(smooth_max_calibration_calls) is not int or smooth_max_calibration_calls < 1: + raise ValueError( + "`smooth_max_calibration_calls` must be a positive integer, " f"got {smooth_max_calibration_calls!r}" + ) + if smooth_eps <= 0: + raise ValueError(f"`smooth_eps` must be positive, got {smooth_eps!r}") + if type(residual_iters) is not int or residual_iters < 1: + raise ValueError(f"`residual_iters` must be a positive integer, got {residual_iters!r}") + if type(residual_early_stop) is not bool: + raise ValueError(f"`residual_early_stop` must be a bool, got {residual_early_stop!r}") + if residual_quant_method != "rtn": + raise ValueError( + "`residual_quant_method` is fixed to 'rtn' by design for SVDQuant residual outer iteration, " + f"got {residual_quant_method!r}" + ) + + self.rank = rank + self.smooth_enabled = smooth_enabled + self.smooth_num_grids = smooth_num_grids + self.smooth_max_calibration_calls = smooth_max_calibration_calls + self.target_modules = _normalize_patterns(target_modules) + self.exclude_modules = _normalize_patterns(exclude_modules) + self.low_rank_dtype = low_rank_dtype + self.smooth_eps = smooth_eps + self.residual_iters = residual_iters + self.residual_early_stop = residual_early_stop + self.residual_quant_method = residual_quant_method + self.model_adapter = model_adapter + self.need_calib = smooth_enabled + + def __repr__(self) -> str: + return ( + f"SVDQuantConfig(rank={self.rank}, smooth_enabled={self.smooth_enabled!r}, " + f"smooth_num_grids={self.smooth_num_grids}, " + f"smooth_max_calibration_calls={self.smooth_max_calibration_calls}, " + f"low_rank_dtype={self.low_rank_dtype!r}, " + f"target_modules={self.target_modules}, exclude_modules={self.exclude_modules}, " + f"residual_iters={self.residual_iters}, residual_early_stop={self.residual_early_stop!r}, " + f"residual_quant_method={self.residual_quant_method!r}, model_adapter={self.model_adapter!r})" + ) + + +def _normalize_patterns(value): + if value is None: + return None + if isinstance(value, str): + return [item.strip() for item in value.split(",") if item.strip()] + return list(value) diff --git a/auto_round/algorithms/transforms/svdquant/residual.py b/auto_round/algorithms/transforms/svdquant/residual.py new file mode 100644 index 0000000000..d23a84c95c --- /dev/null +++ b/auto_round/algorithms/transforms/svdquant/residual.py @@ -0,0 +1,231 @@ +# Copyright (c) 2026 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +from dataclasses import dataclass + +import torch + +from auto_round.data_type.utils import get_quant_func + +_FIXED_MXFP4_DTYPES = frozenset({"mx_fp4", "mx_fp4e2m1"}) +_MXFP4_ALIASES = frozenset({"mx_fp", *_FIXED_MXFP4_DTYPES}) + + +def _validate_scheme_values(scheme): + values = {} + for field in ("data_type", "bits", "group_size", "sym"): + try: + values[field] = getattr(scheme, field) + except AttributeError as exc: + raise ValueError(f"Residual quantization scheme is missing required value {field!r}.") from exc + + if not isinstance(values["data_type"], str) or not values["data_type"].strip(): + raise ValueError("Residual quantization scheme data_type must be a non-empty string.") + if not isinstance(values["bits"], int) or isinstance(values["bits"], bool) or values["bits"] <= 0: + raise ValueError("Residual quantization scheme bits must be a positive integer.") + if values["data_type"] in _FIXED_MXFP4_DTYPES and values["bits"] != 4: + raise ValueError( + f"Residual quantization scheme data_type={values['data_type']!r} requires bits=4; " + f"got bits={values['bits']}." + ) + + group_size = values["group_size"] + scalar_group_size = isinstance(group_size, int) and not isinstance(group_size, bool) and group_size >= -1 + block_group_size = ( + isinstance(group_size, tuple) + and len(group_size) == 2 + and all(isinstance(size, int) and not isinstance(size, bool) and size > 0 for size in group_size) + ) + if not scalar_group_size and not block_group_size: + raise ValueError( + "Residual quantization scheme group_size must be -1, 0, a positive integer, " + "or a pair of positive integers." + ) + if not isinstance(values["sym"], bool): + raise ValueError("Residual quantization scheme sym must be a boolean.") + return values + + +@dataclass(frozen=True) +class ResidualQuantScheme: + """Weight quantization settings for stateless residual QDQ.""" + + data_type: str | None = None + bits: int | None = None + group_size: int | tuple[int, int] | None = None + sym: bool | None = None + + def __post_init__(self) -> None: + _validate_scheme_values(self) + + +@dataclass(frozen=True) +class ActivationQuantScheme: + """Activation quantization settings for stateless calibration QDQ.""" + + data_type: str | None = None + bits: int | None = None + group_size: int | tuple[int, int] | None = None + sym: bool | None = None + + def __post_init__(self) -> None: + _validate_scheme_values(self) + + +@dataclass(frozen=True) +class ResidualDecomposition: + """Best deployment-materialized candidate from residual outer iteration.""" + + residual: torch.Tensor + down: torch.Tensor + up: torch.Tensor + selected_iteration: int + error: float + + +def _rtn_qdq_tensor(tensor: torch.Tensor, scheme, *, tensor_name: str) -> torch.Tensor: + values = _validate_scheme_values(scheme) + requested_dtype = values["data_type"] + if values["bits"] == 4 and requested_dtype in _MXFP4_ALIASES: + requested_dtype = f"{requested_dtype}_rceil" + quant_func, resolved_dtype = get_quant_func( + dtype=requested_dtype, + bits=values["bits"], + sym=values["sym"], + disable_opt_rtn=True, + group_size=values["group_size"], + iters=0, + ) + logical_dtype = resolved_dtype.removeprefix("rtn_") + resolved_base_dtype = logical_dtype.removesuffix("_rceil") + if ( + resolved_base_dtype in _MXFP4_ALIASES + and values["bits"] == 4 + and ( + not isinstance(values["group_size"], int) + or isinstance(values["group_size"], bool) + or values["group_size"] != 32 + ) + ): + raise ValueError( + f"Deployable MXFP4 {tensor_name} QDQ requires scalar group_size=32; " + f"got group_size={values['group_size']!r}." + ) + + qdq, _, _ = quant_func( + tensor=tensor, + bits=values["bits"], + group_size=values["group_size"], + data_type=logical_dtype, + ) + if qdq.shape != tensor.shape or qdq.dtype != tensor.dtype: + raise ValueError( + f"{tensor_name.capitalize()} RTN QDQ must preserve the input shape and dtype; " + f"got shape={tuple(qdq.shape)}, dtype={qdq.dtype}." + ) + if qdq.device != tensor.device: + raise ValueError( + f"{tensor_name.capitalize()} RTN QDQ must preserve the input device; " + f"got input device={tensor.device}, output device={qdq.device}." + ) + if not torch.isfinite(qdq).all(): + raise ValueError(f"{tensor_name.capitalize()} RTN QDQ produced non-finite values.") + return qdq + + +@torch.inference_mode() +def rtn_qdq_residual(weight: torch.Tensor, scheme: ResidualQuantScheme) -> torch.Tensor: + """Apply the registered RTN quantize-dequantize function to a residual.""" + return _rtn_qdq_tensor(weight, scheme, tensor_name="residual") + + +@torch.inference_mode() +def rtn_qdq_activation(activation: torch.Tensor, scheme: ActivationQuantScheme) -> torch.Tensor: + """Apply deployment-compatible dynamic activation quantize-dequantize.""" + return _rtn_qdq_tensor(activation, scheme, tensor_name="activation") + + +@torch.inference_mode() +def truncated_svd(weight: torch.Tensor, rank: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return a rank-limited reconstruction and its shared down/up factors.""" + if weight.ndim != 2: + raise ValueError(f"SVDQuant expects a two-dimensional weight matrix, got shape={tuple(weight.shape)}.") + max_rank = min(weight.shape) + if not isinstance(rank, int) or isinstance(rank, bool) or not 0 <= rank <= max_rank: + raise ValueError(f"SVDQuant rank must be an integer in [0, {max_rank}], got {rank!r}.") + + out_features, in_features = weight.shape + if rank == 0: + low_rank = torch.zeros_like(weight) + down_weight = torch.empty((0, in_features), dtype=weight.dtype, device=weight.device) + up_weight = torch.empty((out_features, 0), dtype=weight.dtype, device=weight.device) + return low_rank, down_weight, up_weight + + u, s, vh = torch.linalg.svd(weight, full_matrices=False) + down_weight = vh[:rank, :] + up_weight = u[:, :rank] * s[:rank].reshape(1, -1) + return up_weight @ down_weight, down_weight, up_weight + + +@torch.inference_mode() +def iterate_residual_decomposition( + weight: torch.Tensor, + *, + rank: int, + scheme: ResidualQuantScheme, + iterations: int, + early_stop: bool, + residual_dtype: torch.dtype, + low_rank_dtype: torch.dtype, +) -> ResidualDecomposition: + """Select the lowest weight-MSE residual/low-rank candidate after deployment casting.""" + if type(iterations) is not int or iterations < 1: + raise ValueError(f"SVDQuant residual iterations must be a positive integer, got {iterations!r}.") + + quantized_residual = torch.zeros_like(weight) + best_down = None + best_up = None + best_error = float("inf") + best_iteration = None + + for iteration in range(1, iterations + 1): + low_rank, down, up = truncated_svd(weight - quantized_residual, rank) + if not all(torch.isfinite(tensor).all() for tensor in (low_rank, down, up)): + break + + residual = (weight - low_rank).to(residual_dtype) + quantized_residual = rtn_qdq_residual(residual, scheme).to(weight.dtype) + deployed_down = down.to(low_rank_dtype) + deployed_up = up.to(low_rank_dtype) + deployed_low_rank = deployed_up.to(weight.dtype) @ deployed_down.to(weight.dtype) + error = torch.sum((weight - (quantized_residual + deployed_low_rank)).square()).item() + accepted = math.isfinite(error) and error <= best_error + + if accepted: + best_down = deployed_down.clone() + best_up = deployed_up.clone() + best_error = error + best_iteration = iteration + elif early_stop and best_iteration is not None: + break + + if best_down is None or best_up is None or best_iteration is None: + raise ValueError("SVDQuant residual iteration did not produce a finite candidate.") + + deployed_low_rank = best_up.to(weight.dtype) @ best_down.to(weight.dtype) + residual = weight - deployed_low_rank + if not torch.isfinite(residual).all(): + raise ValueError("SVDQuant residual iteration produced a non-finite residual.") + return ResidualDecomposition(residual, best_down, best_up, best_iteration, best_error) diff --git a/auto_round/algorithms/transforms/svdquant/wrapper.py b/auto_round/algorithms/transforms/svdquant/wrapper.py new file mode 100644 index 0000000000..7dda281ba8 --- /dev/null +++ b/auto_round/algorithms/transforms/svdquant/wrapper.py @@ -0,0 +1,42 @@ +# Copyright (c) 2026 Intel Corporation +# +# 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. + +from collections.abc import Callable + +import torch + + +class SVDQuantLinear(torch.nn.Module): + """Linear decomposed into a quantizable residual branch and an FP low-rank branch.""" + + def __init__( + self, + residual_linear: torch.nn.Linear, + lora_down: torch.nn.Linear, + lora_up: torch.nn.Linear, + smooth: torch.Tensor, + activation_qdq: Callable[[torch.Tensor], torch.Tensor] | None = None, + ): + super().__init__() + self.residual_linear = residual_linear + self.lora_down = lora_down + self.lora_up = lora_up + self.activation_qdq = activation_qdq + self.register_buffer("smooth", smooth.detach().clone()) + + def forward(self, x): + smooth = self.smooth.to(device=x.device, dtype=x.dtype) + smoothed = x * smooth + residual_input = smoothed if self.activation_qdq is None else self.activation_qdq(smoothed) + return self.residual_linear(residual_input) + self.lora_up(self.lora_down(smoothed)) diff --git a/test/test_cpu/algorithms/test_svdquant_residual.py b/test/test_cpu/algorithms/test_svdquant_residual.py new file mode 100644 index 0000000000..a1dc87dceb --- /dev/null +++ b/test/test_cpu/algorithms/test_svdquant_residual.py @@ -0,0 +1,152 @@ +# Copyright (c) 2026 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch + +import auto_round.algorithms.transforms.svdquant.residual as residual_module +from auto_round.algorithms.transforms.svdquant.config import SVDQuantConfig +from auto_round.algorithms.transforms.svdquant.residual import ( + ResidualQuantScheme, + iterate_residual_decomposition, + rtn_qdq_residual, + truncated_svd, +) +from auto_round.algorithms.transforms.svdquant.wrapper import SVDQuantLinear +from auto_round.data_type.mxfp import quant_mx_rceil + + +def test_rtn_qdq_residual_matches_deployable_mxfp4_quantizer(): + weight = torch.linspace(-3.0, 3.0, steps=3 * 64, dtype=torch.float32).reshape(3, 64) + scheme = ResidualQuantScheme(data_type="mx_fp4e2m1", bits=4, group_size=32, sym=True) + + actual = rtn_qdq_residual(weight, scheme) + expected, _, _ = quant_mx_rceil( + tensor=weight, + bits=4, + group_size=32, + data_type="mx_fp4e2m1", + ) + + torch.testing.assert_close(actual, expected) + assert actual.shape == weight.shape + assert actual.dtype == weight.dtype + assert actual.device == weight.device + + +def test_svdquant_linear_combines_residual_and_low_rank_branches_after_smoothing(): + residual = torch.nn.Linear(4, 3, bias=True) + lora_down = torch.nn.Linear(4, 2, bias=False) + lora_up = torch.nn.Linear(2, 3, bias=False) + smooth = torch.tensor([0.5, 1.0, 2.0, 4.0]) + wrapper = SVDQuantLinear(residual, lora_down, lora_up, smooth) + inputs = torch.randn(2, 4) + + smoothed = inputs * smooth + expected = residual(smoothed) + lora_up(lora_down(smoothed)) + + torch.testing.assert_close(wrapper(inputs), expected) + + +def test_svdquant_config_defaults_to_data_free_single_iteration(): + config = SVDQuantConfig() + + assert config.rank == 32 + assert config.smooth_enabled is False + assert config.smooth_max_calibration_calls == 128 + assert config.residual_iters == 1 + assert config.residual_early_stop is False + assert config.need_calib is False + + +@pytest.mark.parametrize( + ("kwargs", "field"), + [ + ({"rank": -1}, "rank"), + ({"smooth_enabled": 1}, "smooth_enabled"), + ({"smooth_num_grids": 1}, "smooth_num_grids"), + ({"smooth_max_calibration_calls": 0}, "smooth_max_calibration_calls"), + ({"residual_iters": 0}, "residual_iters"), + ({"residual_quant_method": "signround"}, "residual_quant_method"), + ], +) +def test_svdquant_config_rejects_invalid_structural_options(kwargs, field): + with pytest.raises(ValueError, match=field): + SVDQuantConfig(**kwargs) + + +def test_truncated_svd_returns_shared_down_factor_for_stacked_projection_group(): + torch.manual_seed(0) + qkv = torch.randn(12, 8, dtype=torch.float32) + + low_rank, down, up = truncated_svd(qkv, rank=3) + + assert low_rank.shape == qkv.shape + assert down.shape == (3, 8) + assert up.shape == (12, 3) + torch.testing.assert_close(low_rank, up @ down) + q_up, k_up, v_up = up.split((4, 4, 4), dim=0) + torch.testing.assert_close(torch.cat((q_up @ down, k_up @ down, v_up @ down)), low_rank) + + +def test_residual_iteration_keeps_the_best_materialized_candidate(): + torch.manual_seed(1) + weight = torch.randn(8, 32, dtype=torch.float32) + scheme = ResidualQuantScheme(data_type="mx_fp4e2m1", bits=4, group_size=32, sym=True) + + first_low_rank, first_down, first_up = truncated_svd(weight, rank=2) + first_qdq = rtn_qdq_residual(weight - first_low_rank, scheme) + first_error = torch.sum((weight - (first_qdq + first_up @ first_down)).square()).item() + + result = iterate_residual_decomposition( + weight, + rank=2, + scheme=scheme, + iterations=4, + early_stop=False, + residual_dtype=torch.float32, + low_rank_dtype=torch.float32, + ) + + assert 1 <= result.selected_iteration <= 4 + assert result.error <= first_error + assert result.residual.shape == weight.shape + torch.testing.assert_close(result.residual + result.up @ result.down, weight) + + +def test_residual_iteration_early_stops_after_candidate_worsens(monkeypatch): + torch.manual_seed(2) + weight = torch.randn(4, 32, dtype=torch.float32) + scheme = ResidualQuantScheme(data_type="mx_fp4e2m1", bits=4, group_size=32, sym=True) + calls = 0 + + def qdq_then_degrade(residual, _scheme): + nonlocal calls + calls += 1 + return residual if calls == 1 else residual + 1 + + monkeypatch.setattr(residual_module, "rtn_qdq_residual", qdq_then_degrade) + + result = iterate_residual_decomposition( + weight, + rank=2, + scheme=scheme, + iterations=10, + early_stop=True, + residual_dtype=torch.float32, + low_rank_dtype=torch.float32, + ) + + assert calls == 2 + assert result.selected_iteration == 1 From 012ffaa84c7bf9a6bda97c5addf68a4496365b12 Mon Sep 17 00:00:00 2001 From: changwangss Date: Tue, 28 Jul 2026 16:22:13 +0800 Subject: [PATCH 2/9] feat: integrate SVDQuant with algorithm composer Signed-off-by: changwangss --- auto_round/algorithms/registry.py | 1 + .../transforms/svdquant/__init__.py | 3 +- .../algorithms/transforms/svdquant/apply.py | 291 ++++++++++++++++++ .../svdquant/smooth_adapters/__init__.py | 34 ++ .../svdquant/smooth_adapters/base.py | 93 ++++++ .../svdquant/smooth_adapters/flux.py | 93 ++++++ test/test_cpu/algorithms/test_svdquant.py | 131 ++++++++ 7 files changed, 645 insertions(+), 1 deletion(-) create mode 100644 auto_round/algorithms/transforms/svdquant/apply.py create mode 100644 auto_round/algorithms/transforms/svdquant/smooth_adapters/__init__.py create mode 100644 auto_round/algorithms/transforms/svdquant/smooth_adapters/base.py create mode 100644 auto_round/algorithms/transforms/svdquant/smooth_adapters/flux.py create mode 100644 test/test_cpu/algorithms/test_svdquant.py diff --git a/auto_round/algorithms/registry.py b/auto_round/algorithms/registry.py index d56b0a4970..4380a0912a 100644 --- a/auto_round/algorithms/registry.py +++ b/auto_round/algorithms/registry.py @@ -47,6 +47,7 @@ def _ensure_pipeline_members_registered() -> None: "auto_round.algorithms.quantization.sign_roundv2.quantizer", "auto_round.algorithms.quantization.adam_round.adam", "auto_round.algorithms.transforms.awq.base", + "auto_round.algorithms.transforms.svdquant.apply", ): importlib.import_module(module_name) _pipeline_members_registered = True diff --git a/auto_round/algorithms/transforms/svdquant/__init__.py b/auto_round/algorithms/transforms/svdquant/__init__.py index 14099b8a10..ea42781f89 100644 --- a/auto_round/algorithms/transforms/svdquant/__init__.py +++ b/auto_round/algorithms/transforms/svdquant/__init__.py @@ -12,7 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +from auto_round.algorithms.transforms.svdquant.apply import SVDQuantTransform from auto_round.algorithms.transforms.svdquant.config import SVDQuantConfig from auto_round.algorithms.transforms.svdquant.wrapper import SVDQuantLinear -__all__ = ["SVDQuantConfig", "SVDQuantLinear"] +__all__ = ["SVDQuantConfig", "SVDQuantLinear", "SVDQuantTransform"] diff --git a/auto_round/algorithms/transforms/svdquant/apply.py b/auto_round/algorithms/transforms/svdquant/apply.py new file mode 100644 index 0000000000..79116e2dd9 --- /dev/null +++ b/auto_round/algorithms/transforms/svdquant/apply.py @@ -0,0 +1,291 @@ +# Copyright (c) 2026 Intel Corporation +# +# 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. + +from __future__ import annotations + +import torch + +from auto_round.algorithms.registry import register_pipeline_member +from auto_round.algorithms.transforms.base import BasePreprocessor +from auto_round.algorithms.transforms.svdquant.config import SVDQuantConfig +from auto_round.algorithms.transforms.svdquant.residual import ( + ResidualQuantScheme, + iterate_residual_decomposition, + truncated_svd, +) +from auto_round.algorithms.transforms.svdquant.smooth_adapters import SmoothSearchGroup, discover_svdquant_groups +from auto_round.algorithms.transforms.svdquant.wrapper import SVDQuantLinear +from auto_round.logger import logger + +_SCHEME_ATTRS = ( + "bits", + "group_size", + "sym", + "data_type", + "act_bits", + "act_group_size", + "act_sym", + "act_data_type", + "act_dynamic", + "super_bits", + "super_group_size", + "super_sym", + "scale_dtype", + "weight_global_scale", + "tuning_device", +) +_MXFP4_ALIASES = frozenset({"mx_fp", "mx_fp4", "mx_fp4e2m1"}) + + +@register_pipeline_member(SVDQuantConfig) +class SVDQuantTransform(BasePreprocessor): + """Split target linears into a quantized residual and an FP low-rank branch.""" + + def __init__(self, config: SVDQuantConfig) -> None: + super().__init__(config) + self._configured_block_names: tuple[str, ...] = () + self._block_groups: dict[str, list[SmoothSearchGroup]] = {} + + def bind(self, orchestrator) -> None: + super().bind(orchestrator) + nblocks = getattr(orchestrator, "nblocks", 1) + if nblocks != 1: + raise ValueError(f"SVDQuant requires nblocks=1, got nblocks={nblocks}.") + quant_block_list = getattr(orchestrator, "quant_block_list", None) or () + self._configured_block_names = tuple( + block_name for block_group in quant_block_list for block_name in block_group + ) + + def prepare_run(self, composer=None) -> None: + self._block_groups.clear() + if self.model is None: + return + for block_name in self._configured_block_names: + block = self.model.get_submodule(block_name) + self._block_groups[block_name] = discover_svdquant_groups(block, self._is_target) + logger.info( + "SVDQuant: resolved %d projection groups across %d blocks.", + sum(len(groups) for groups in self._block_groups.values()), + len(self._block_groups), + ) + + def register_fp_input_forward_hooks(self, block) -> list: + if self.config.smooth_enabled: + raise NotImplementedError("SVDQuant smooth calibration is not ported to the main architecture yet.") + return [] + + @torch.no_grad() + def pre_quantize_block(self, ctx) -> None: + if len(ctx.block_names) != 1: + raise ValueError(f"SVDQuant requires one block at a time, got {ctx.block_names!r}.") + if self.config.smooth_enabled: + raise NotImplementedError("SVDQuant smooth calibration is not ported to the main architecture yet.") + + block_name = ctx.block_name + block = ctx.model.get_submodule(block_name) + groups = self._block_groups.get(block_name) + if groups is None: + groups = discover_svdquant_groups(block, self._is_target) + self._block_groups[block_name] = groups + + local_names = {id(module): name for name, module in block.named_modules() if name} + replacements = [] + for group in groups: + wrappers = self._decompose_group(group) + for projection, wrapper in zip(group.projections, wrappers): + local_name = local_names.get(id(projection)) + if local_name is None: + raise ValueError(f"SVDQuant could not locate projection {self._module_name(projection)!r}.") + replacements.append((local_name, wrapper)) + + for local_name, wrapper in replacements: + _set_child_module(block, local_name, wrapper) + + def post_quantize_block(self, ctx) -> None: + self._block_groups.pop(ctx.block_name, None) + + def finalize_run(self) -> None: + self._block_groups.clear() + + def _decompose_group(self, group: SmoothSearchGroup) -> list[SVDQuantLinear]: + weights = [projection.weight.detach().to(torch.float32) for projection in group.projections] + stacked = torch.cat(weights, dim=0) + output_sizes = [projection.out_features for projection in group.projections] + rank = min(self.config.rank, *stacked.shape) + low_rank_dtype = self._resolve_low_rank_dtype(group.projections[0].weight.dtype) + + if self.config.residual_iters == 1: + _, down, up = truncated_svd(stacked, rank) + deployed_down = down.to(low_rank_dtype) + deployed_up = up.to(low_rank_dtype) + deployed_low_rank = deployed_up.float() @ deployed_down.float() + residuals = [ + residual.to(projection.weight.dtype) + for residual, projection in zip( + (stacked - deployed_low_rank).split(output_sizes, dim=0), group.projections + ) + ] + else: + scheme = self._shared_residual_scheme(group) + result = iterate_residual_decomposition( + stacked, + rank=rank, + scheme=scheme, + iterations=self.config.residual_iters, + early_stop=self.config.residual_early_stop, + residual_dtype=group.projections[0].weight.dtype, + low_rank_dtype=low_rank_dtype, + ) + deployed_down = result.down + deployed_up = result.up + residuals = [ + residual.to(projection.weight.dtype) + for residual, projection in zip(result.residual.split(output_sizes, dim=0), group.projections) + ] + logger.info( + "SVDQuant residual selected iteration %d/%d for %s: weight error %.6g", + result.selected_iteration, + self.config.residual_iters, + group.key, + result.error, + ) + + up_parts = deployed_up.split(output_sizes, dim=0) + smooth = torch.ones(stacked.shape[1], device=stacked.device, dtype=torch.float32) + return self._build_group_wrappers(group, residuals, deployed_down, up_parts, smooth) + + def _build_group_wrappers( + self, + group: SmoothSearchGroup, + residuals: list[torch.Tensor], + down: torch.Tensor, + up_parts: tuple[torch.Tensor, ...], + smooth: torch.Tensor, + ) -> list[SVDQuantLinear]: + wrappers = [] + rank = down.shape[0] + for projection, residual_weight, up in zip(group.projections, residuals, up_parts): + residual = self._new_linear_like(projection, residual_weight, projection.bias) + lora_down = torch.nn.Linear( + projection.in_features, + rank, + bias=False, + dtype=down.dtype, + device=projection.weight.device, + ) + lora_up = torch.nn.Linear( + rank, + projection.out_features, + bias=False, + dtype=up.dtype, + device=projection.weight.device, + ) + lora_down.weight.copy_(down) + lora_up.weight.copy_(up) + self._mark_unquantized(lora_down) + self._mark_unquantized(lora_up) + self._copy_quant_attrs(projection, residual, suffix=".residual_linear") + wrappers.append(SVDQuantLinear(residual, lora_down, lora_up, smooth.to(projection.weight.dtype))) + return wrappers + + def _shared_residual_scheme(self, group: SmoothSearchGroup) -> ResidualQuantScheme: + schemes = tuple(self._residual_quant_scheme(projection) for projection in group.projections) + if any(scheme != schemes[0] for scheme in schemes[1:]): + raise ValueError(f"SVDQuant group {group.key!r} has inconsistent residual quantization schemes.") + return schemes[0] + + def _residual_quant_scheme(self, module: torch.nn.Linear) -> ResidualQuantScheme: + required = ("data_type", "bits", "group_size", "sym") + missing = [attr for attr in required if not hasattr(module, attr) or getattr(module, attr) is None] + if missing: + raise ValueError( + f"SVDQuant residual iteration requires a complete quantization scheme for " + f"{self._module_name(module)!r}; missing: {', '.join(missing)}." + ) + return ResidualQuantScheme(**{attr: getattr(module, attr) for attr in required}) + + def _is_target(self, name: str, module: torch.nn.Module) -> bool: + if not isinstance(module, torch.nn.Linear): + return False + full_name = str(getattr(module, "global_name", name)) + if self.config.target_modules and not any( + pattern in name or pattern in full_name for pattern in self.config.target_modules + ): + return False + if self.config.exclude_modules and any( + pattern in name or pattern in full_name for pattern in self.config.exclude_modules + ): + return False + return True + + @staticmethod + def _new_linear_like(module: torch.nn.Linear, weight: torch.Tensor, bias: torch.Tensor | None): + residual = torch.nn.Linear( + module.in_features, + module.out_features, + bias=bias is not None, + dtype=module.weight.dtype, + device=module.weight.device, + ) + residual.weight.copy_(weight.to(module.weight.dtype)) + if bias is not None: + residual.bias.copy_(bias.detach().to(module.weight.dtype)) + return residual + + @staticmethod + def _mark_unquantized(module: torch.nn.Module) -> None: + module.bits = 16 + module.act_bits = 16 + + @staticmethod + def _copy_quant_attrs(src: torch.nn.Module, dst: torch.nn.Module, suffix: str) -> None: + for attr in _SCHEME_ATTRS: + if hasattr(src, attr): + setattr(dst, attr, getattr(src, attr)) + if getattr(dst, "bits", None) == 4 and getattr(dst, "data_type", None) in _MXFP4_ALIASES: + dst.data_type = f"{dst.data_type}_rceil" + if getattr(dst, "act_bits", None) == 4 and getattr(dst, "act_data_type", None) in _MXFP4_ALIASES: + dst.act_data_type = f"{dst.act_data_type}_rceil" + if hasattr(src, "global_name"): + dst.global_name = f"{src.global_name}{suffix}" + + def _resolve_low_rank_dtype(self, fallback: torch.dtype) -> torch.dtype: + dtype = str(self.config.low_rank_dtype).lower() + if dtype in {"bf16", "bfloat16"}: + return torch.bfloat16 + if dtype in {"fp16", "float16"}: + return torch.float16 + if dtype in {"fp32", "float32"}: + return torch.float32 + return fallback + + @staticmethod + def _module_name(module: torch.nn.Module) -> str: + return str(getattr(module, "global_name", module.__class__.__name__)) + + +def _set_child_module(root: torch.nn.Module, name: str, module: torch.nn.Module) -> None: + parts = name.split(".") + parent = root + for part in parts[:-1]: + parent = ( + parent[int(part)] + if part.isdigit() and isinstance(parent, (torch.nn.ModuleList, torch.nn.Sequential)) + else getattr(parent, part) + ) + leaf = parts[-1] + if leaf.isdigit() and isinstance(parent, (torch.nn.ModuleList, torch.nn.Sequential)): + parent[int(leaf)] = module + else: + setattr(parent, leaf, module) diff --git a/auto_round/algorithms/transforms/svdquant/smooth_adapters/__init__.py b/auto_round/algorithms/transforms/svdquant/smooth_adapters/__init__.py new file mode 100644 index 0000000000..b79f0b0e41 --- /dev/null +++ b/auto_round/algorithms/transforms/svdquant/smooth_adapters/__init__.py @@ -0,0 +1,34 @@ +# Copyright (c) 2026 Intel Corporation +# +# 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. + +from __future__ import annotations + +import torch + +from auto_round.algorithms.transforms.svdquant.smooth_adapters.base import ( + SmoothSearchGroup, + TargetPredicate, + generic_linear_groups, +) +from auto_round.algorithms.transforms.svdquant.smooth_adapters.flux import discover_flux_groups, supports_flux_block + + +def discover_svdquant_groups(block: torch.nn.Module, is_target: TargetPredicate) -> list[SmoothSearchGroup]: + """Discover shared-input projection groups for one quantization block.""" + if supports_flux_block(block): + return discover_flux_groups(block, is_target) + return generic_linear_groups(block, is_target) + + +__all__ = ["SmoothSearchGroup", "discover_svdquant_groups"] diff --git a/auto_round/algorithms/transforms/svdquant/smooth_adapters/base.py b/auto_round/algorithms/transforms/svdquant/smooth_adapters/base.py new file mode 100644 index 0000000000..34a9179407 --- /dev/null +++ b/auto_round/algorithms/transforms/svdquant/smooth_adapters/base.py @@ -0,0 +1,93 @@ +# Copyright (c) 2026 Intel Corporation +# +# 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. + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass + +import torch + +TargetPredicate = Callable[[str, torch.nn.Module], bool] + + +@dataclass(frozen=True) +class SmoothSearchGroup: + """Projection group that shares an input scale and low-rank down factor.""" + + key: str + projection_names: tuple[str, ...] + projections: tuple[torch.nn.Linear, ...] + projection_input_module: torch.nn.Module + evaluation_module: torch.nn.Module + output_indices: tuple[int, ...] | None = None + output_splits: tuple[int, ...] = () + + def __post_init__(self) -> None: + if not self.projections: + raise ValueError(f"SVDQuant group {self.key!r} has no projections.") + if len(self.projection_names) != len(self.projections): + raise ValueError(f"SVDQuant group {self.key!r} has mismatched names and projections.") + if len({projection.in_features for projection in self.projections}) != 1: + raise ValueError(f"SVDQuant group {self.key!r} projections must share an input width.") + if self.output_splits and sum(self.output_splits) != len(self.projections): + raise ValueError(f"SVDQuant group {self.key!r} output splits do not cover its projections.") + + +def module_global_name(block: torch.nn.Module, local_name: str) -> str: + prefix = str(getattr(block, "global_name", block.__class__.__name__)) + return f"{prefix}.{local_name}" if local_name else prefix + + +def resolve_module(root: torch.nn.Module, path: str) -> torch.nn.Module | None: + module = root + for part in path.split("."): + if part.isdigit() and isinstance(module, (torch.nn.ModuleList, torch.nn.Sequential)): + index = int(part) + if index >= len(module): + return None + module = module[index] + else: + candidate = getattr(module, part, None) + if not isinstance(candidate, torch.nn.Module): + return None + module = candidate + return module + + +def generic_linear_groups( + block: torch.nn.Module, + is_target: TargetPredicate, + *, + consumed: set[int] | None = None, +) -> list[SmoothSearchGroup]: + consumed = consumed or set() + groups = [] + for local_name, module in block.named_modules(): + if not local_name or id(module) in consumed or not is_target(local_name, module): + continue + if not isinstance(module, torch.nn.Linear): + continue + name = module_global_name(block, local_name) + groups.append( + SmoothSearchGroup( + key=name, + projection_names=(name,), + projections=(module,), + projection_input_module=module, + evaluation_module=module, + output_splits=(1,), + ) + ) + return groups diff --git a/auto_round/algorithms/transforms/svdquant/smooth_adapters/flux.py b/auto_round/algorithms/transforms/svdquant/smooth_adapters/flux.py new file mode 100644 index 0000000000..e7b4d5d30f --- /dev/null +++ b/auto_round/algorithms/transforms/svdquant/smooth_adapters/flux.py @@ -0,0 +1,93 @@ +# Copyright (c) 2026 Intel Corporation +# +# 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. + +from __future__ import annotations + +import torch + +from auto_round.algorithms.transforms.svdquant.smooth_adapters.base import ( + SmoothSearchGroup, + TargetPredicate, + generic_linear_groups, + module_global_name, + resolve_module, +) + +_DOUBLE_GROUPS = ( + ("attn.qkv", ("attn.to_q", "attn.to_k", "attn.to_v"), "attn", (0,), (3,)), + ("attn.add_qkv", ("attn.add_q_proj", "attn.add_k_proj", "attn.add_v_proj"), "attn", (1,), (3,)), + ("attn.to_out.0", ("attn.to_out.0",), "attn.to_out.0", None, (1,)), + ("attn.to_add_out", ("attn.to_add_out",), "attn.to_add_out", None, (1,)), + ("ff.net.0.proj", ("ff.net.0.proj",), "ff.net.0.proj", None, (1,)), + ("ff.net.2", ("ff.net.2",), "ff.net.2", None, (1,)), + ("ff_context.net.0.proj", ("ff_context.net.0.proj",), "ff_context.net.0.proj", None, (1,)), + ("ff_context.net.2", ("ff_context.net.2",), "ff_context.net.2", None, (1,)), +) + + +def supports_flux_block(block: torch.nn.Module) -> bool: + return block.__class__.__name__ in {"FluxTransformerBlock", "FluxSingleTransformerBlock"} + + +def _make_group( + block: torch.nn.Module, + key: str, + projection_paths: tuple[str, ...], + evaluation_path: str, + output_indices: tuple[int, ...] | None, + output_splits: tuple[int, ...], + is_target: TargetPredicate, +) -> SmoothSearchGroup | None: + selected = [] + for path in projection_paths: + module = resolve_module(block, path) + if isinstance(module, torch.nn.Linear) and is_target(path, module): + selected.append((path, module)) + if not selected: + return None + + evaluation_module = block if not evaluation_path else resolve_module(block, evaluation_path) + if evaluation_module is None: + raise ValueError(f"Flux SVDQuant group {module_global_name(block, key)!r} has no evaluation module.") + names = tuple(module_global_name(block, path) for path, _ in selected) + projections = tuple(module for _, module in selected) + splits = output_splits if sum(output_splits) == len(projections) else (len(projections),) + return SmoothSearchGroup( + key=module_global_name(block, key), + projection_names=names, + projections=projections, + projection_input_module=projections[0], + evaluation_module=evaluation_module, + output_indices=output_indices, + output_splits=splits, + ) + + +def discover_flux_groups(block: torch.nn.Module, is_target: TargetPredicate) -> list[SmoothSearchGroup]: + if block.__class__.__name__ == "FluxSingleTransformerBlock": + specifications = ( + ("parallel_qkv_mlp", ("attn.to_q", "attn.to_k", "attn.to_v", "proj_mlp"), "", None, (3, 1)), + ("proj_out", ("proj_out",), "proj_out", None, (1,)), + ) + else: + specifications = _DOUBLE_GROUPS + + groups = [] + for specification in specifications: + group = _make_group(block, *specification, is_target) + if group is not None: + groups.append(group) + consumed = {id(projection) for group in groups for projection in group.projections} + groups.extend(generic_linear_groups(block, is_target, consumed=consumed)) + return groups diff --git a/test/test_cpu/algorithms/test_svdquant.py b/test/test_cpu/algorithms/test_svdquant.py new file mode 100644 index 0000000000..7b010a1408 --- /dev/null +++ b/test/test_cpu/algorithms/test_svdquant.py @@ -0,0 +1,131 @@ +# Copyright (c) 2026 Intel Corporation +# +# 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. + +from types import SimpleNamespace + +import torch + +from auto_round.algorithms.composer import AlgorithmComposer, BlockContext +from auto_round.algorithms.quantization.rtn.config import RTNConfig +from auto_round.algorithms.transforms.svdquant import SVDQuantConfig, SVDQuantLinear +from auto_round.algorithms.transforms.svdquant.apply import SVDQuantTransform + + +class FluxAttention(torch.nn.Module): + def __init__(self, width=8): + super().__init__() + self.to_q = torch.nn.Linear(width, 4, bias=False) + self.to_k = torch.nn.Linear(width, 4, bias=False) + self.to_v = torch.nn.Linear(width, 4, bias=False) + + +class FluxTransformerBlock(torch.nn.Module): + def __init__(self, width=8): + super().__init__() + self.attn = FluxAttention(width) + + +class TinyFlux(torch.nn.Module): + def __init__(self, width=8): + super().__init__() + self.transformer_blocks = torch.nn.ModuleList([FluxTransformerBlock(width)]) + + +def _mark_modules(model): + for name, module in model.named_modules(): + module.global_name = name + if isinstance(module, torch.nn.Linear): + module.bits = 4 + module.group_size = 32 + module.sym = True + module.data_type = "mx_fp4e2m1" + module.act_bits = 16 + + +def test_svdquant_composes_before_rtn(): + composer = AlgorithmComposer([SVDQuantConfig(rank=2), RTNConfig(disable_opt_rtn=True)]) + + assert len(composer.preprocessors) == 1 + assert isinstance(composer.preprocessors[0], SVDQuantTransform) + + +def test_no_smooth_flux_qkv_share_one_down_factor(): + model = TinyFlux() + _mark_modules(model) + block_name = "transformer_blocks.0" + inputs = torch.randn(3, 8) + expected = tuple( + projection(inputs) + for projection in ( + model.transformer_blocks[0].attn.to_q, + model.transformer_blocks[0].attn.to_k, + model.transformer_blocks[0].attn.to_v, + ) + ) + transform = SVDQuantTransform(SVDQuantConfig(rank=2, residual_iters=1, low_rank_dtype="fp32")) + orchestrator = SimpleNamespace( + model_context=SimpleNamespace(model=model), + compress_context=None, + calibration_context=None, + scheme_context=None, + scale_dtype=None, + nblocks=1, + quant_block_list=[[block_name]], + ) + transform.bind(orchestrator) + transform.prepare_run() + + transform.pre_quantize_block( + BlockContext(model=model, block_names=[block_name], block_name=block_name, block_index=0) + ) + + q = model.transformer_blocks[0].attn.to_q + k = model.transformer_blocks[0].attn.to_k + v = model.transformer_blocks[0].attn.to_v + assert all(isinstance(module, SVDQuantLinear) for module in (q, k, v)) + torch.testing.assert_close(q.lora_down.weight, k.lora_down.weight) + torch.testing.assert_close(q.lora_down.weight, v.lora_down.weight) + assert q.residual_linear.data_type == "mx_fp4e2m1_rceil" + assert q.lora_down.bits == 16 + assert q.lora_up.bits == 16 + for actual, reference in zip((q(inputs), k(inputs), v(inputs)), expected): + torch.testing.assert_close(actual, reference) + + +def test_no_smooth_grouped_residual_iteration_and_cleanup(): + model = TinyFlux(width=32) + _mark_modules(model) + block_name = "transformer_blocks.0" + transform = SVDQuantTransform( + SVDQuantConfig(rank=2, residual_iters=2, residual_early_stop=True, low_rank_dtype="fp32") + ) + orchestrator = SimpleNamespace( + model_context=SimpleNamespace(model=model), + compress_context=None, + calibration_context=None, + scheme_context=None, + scale_dtype=None, + nblocks=1, + quant_block_list=[[block_name]], + ) + transform.bind(orchestrator) + transform.prepare_run() + ctx = BlockContext(model=model, block_names=[block_name], block_name=block_name, block_index=0) + + transform.pre_quantize_block(ctx) + + assert isinstance(model.transformer_blocks[0].attn.to_q, SVDQuantLinear) + assert block_name in transform._block_groups + transform.post_quantize_block(ctx) + assert block_name not in transform._block_groups From 493cf94a9b3c5fb0446ff36cb1bc0d51c3088ef1 Mon Sep 17 00:00:00 2001 From: changwangss Date: Tue, 28 Jul 2026 16:28:13 +0800 Subject: [PATCH 3/9] feat: add SVDQuant algorithm CLI Signed-off-by: changwangss --- auto_round/cli/algorithms.py | 98 +++++++++++++++++++++++++++ test/test_cpu/utils/test_cli_usage.py | 81 ++++++++++++++++++++++ 2 files changed, 179 insertions(+) diff --git a/auto_round/cli/algorithms.py b/auto_round/cli/algorithms.py index cb18306e53..718e23c7ba 100644 --- a/auto_round/cli/algorithms.py +++ b/auto_round/cli/algorithms.py @@ -230,6 +230,96 @@ def build(self, args, common_kwargs: dict[str, Any]): ) +class SVDQuant(AlgorithmHandler): + name = "svdquant" + aliases = ("svdquant",) + summary = "SVD low-rank decomposition before residual quantization." + config_factory = None + + def register(self, group) -> None: + group.add_argument("--svdquant-rank", default=32, type=int, help="SVDQuant low-rank size.") + group.add_argument( + "--enable-svdquant-smooth", + dest="svdquant_smooth_enabled", + default=False, + action="store_true", + help="Enable SVDQuant activation-aware smoothing.", + ) + group.add_argument( + "--svdquant-smooth-num-grids", + default=20, + type=int, + help="Number of candidates per SVDQuant smooth search grid family.", + ) + group.add_argument( + "--svdquant-smooth-max-calibration-calls", + default=128, + type=int, + help="Maximum calibration calls retained per SVDQuant smooth group.", + ) + group.add_argument( + "--svdquant-residual-iters", + default=1, + type=int, + help="Number of alternating low-rank and residual quantization iterations.", + ) + group.add_argument( + "--enable-svdquant-residual-early-stop", + dest="svdquant_residual_early_stop", + default=False, + action="store_true", + help="Stop residual iteration when reconstruction error no longer improves.", + ) + group.add_argument( + "--svdquant-residual-quant-method", + default="rtn", + choices=["rtn"], + help="Residual outer iteration quantizer; fixed to RTN independently of the terminal quantizer.", + ) + group.add_argument( + "--svdquant-low-rank-dtype", + default="bf16", + choices=["bf16", "bfloat16", "fp16", "float16", "fp32", "float32"], + help="Data type for the SVDQuant low-rank branch.", + ) + group.add_argument( + "--svdquant-target-modules", + default=None, + type=str, + help="Comma-separated module-name substrings to transform.", + ) + group.add_argument( + "--svdquant-exclude-modules", + default=None, + type=str, + help="Comma-separated module-name substrings to exclude.", + ) + group.add_argument( + "--svdquant-model-adapter", + default="auto", + choices=["auto", "identity", "flux"], + help="Architecture adapter used by SVDQuant export.", + ) + + def build(self, args, common_kwargs: dict[str, Any]): + from auto_round.algorithms.transforms.svdquant.config import SVDQuantConfig + + return SVDQuantConfig( + rank=getattr(args, "svdquant_rank", 32), + smooth_enabled=getattr(args, "svdquant_smooth_enabled", False), + smooth_num_grids=getattr(args, "svdquant_smooth_num_grids", 20), + smooth_max_calibration_calls=getattr(args, "svdquant_smooth_max_calibration_calls", 128), + residual_iters=getattr(args, "svdquant_residual_iters", 1), + residual_early_stop=getattr(args, "svdquant_residual_early_stop", False), + residual_quant_method=getattr(args, "svdquant_residual_quant_method", "rtn"), + low_rank_dtype=getattr(args, "svdquant_low_rank_dtype", "bf16"), + target_modules=getattr(args, "svdquant_target_modules", None), + exclude_modules=getattr(args, "svdquant_exclude_modules", None), + model_adapter=getattr(args, "svdquant_model_adapter", "auto"), + **common_kwargs, + ) + + class RTN(AlgorithmHandler): name = "rtn" aliases = ("rtn",) @@ -415,6 +505,7 @@ def _register_builtin_algorithm_factories() -> None: from auto_round.algorithms.quantization.sign_round.config import SignRoundConfig from auto_round.algorithms.transforms.awq.config import AWQConfig from auto_round.algorithms.transforms.hadamard.config import RotationConfig + from auto_round.algorithms.transforms.svdquant.config import SVDQuantConfig register_algorithm("rtn", aliases=("rtn",), config_factory=RTNConfig, cli_handler=RTN, summary=RTN.summary) register_algorithm( @@ -425,6 +516,13 @@ def _register_builtin_algorithm_factories() -> None: summary=AutoRound.summary, ) register_algorithm("awq", aliases=("awq",), config_factory=AWQConfig, cli_handler=AWQ, summary=AWQ.summary) + register_algorithm( + "svdquant", + aliases=("svdquant",), + config_factory=SVDQuantConfig, + cli_handler=SVDQuant, + summary=SVDQuant.summary, + ) register_algorithm( "hadamard", aliases=("hadamard", "random_hadamard", "quarot_hadamard"), diff --git a/test/test_cpu/utils/test_cli_usage.py b/test/test_cpu/utils/test_cli_usage.py index ac937ec92b..2df5c5f4a5 100644 --- a/test/test_cpu/utils/test_cli_usage.py +++ b/test/test_cpu/utils/test_cli_usage.py @@ -251,6 +251,87 @@ def test_legacy_disable_flags_map_to_enable_bools(): assert args.enable_quanted_input is False +def test_svdquant_cli_builds_hyphenated_options_before_rtn(): + from auto_round.algorithms.transforms.svdquant.config import SVDQuantConfig + from auto_round.cli.algorithms import AlgorithmHandler + from auto_round.cli.parser import build_quantize_parser + + parser = build_quantize_parser() + args = parser.parse_args( + [ + "--model", + "dummy-model", + "--algorithm", + "svdquant,rtn", + "--svdquant-rank", + "16", + "--enable-svdquant-smooth", + "--svdquant-smooth-num-grids", + "39", + "--svdquant-smooth-max-calibration-calls", + "64", + "--svdquant-residual-iters", + "20", + "--enable-svdquant-residual-early-stop", + "--svdquant-residual-quant-method", + "rtn", + "--svdquant-low-rank-dtype", + "fp32", + "--svdquant-target-modules", + "attn,ff", + "--svdquant-exclude-modules", + "proj_out", + "--svdquant-model-adapter", + "flux", + "--disable_opt_rtn", + ] + ) + + configs = AlgorithmHandler.build_configs(args, {}) + + assert isinstance(configs[0], SVDQuantConfig) + assert configs[0].rank == 16 + assert configs[0].smooth_enabled is True + assert configs[0].smooth_num_grids == 39 + assert configs[0].smooth_max_calibration_calls == 64 + assert configs[0].residual_iters == 20 + assert configs[0].residual_early_stop is True + assert configs[0].low_rank_dtype == "fp32" + assert configs[0].target_modules == ["attn", "ff"] + assert configs[0].exclude_modules == ["proj_out"] + assert configs[0].model_adapter == "flux" + assert configs[1].__class__.__name__ == "RTNConfig" + + +def test_svdquant_cli_rejects_underscore_option_aliases(): + from auto_round.cli.parser import build_quantize_parser + + with pytest.raises(SystemExit): + build_quantize_parser().parse_args( + ["--model", "dummy-model", "--algorithm", "svdquant,rtn", "--svdquant_rank", "16"] + ) + + +def test_svdquant_cli_defaults_compose_before_signround(): + from auto_round.algorithms.quantization.sign_round.config import SignRoundConfig + from auto_round.algorithms.transforms.svdquant.config import SVDQuantConfig + from auto_round.cli.algorithms import AlgorithmHandler + from auto_round.cli.parser import build_quantize_parser + + args = build_quantize_parser().parse_args( + ["--model", "dummy-model", "--algorithm", "svdquant,auto_round", "--iters", "200"] + ) + + configs = AlgorithmHandler.build_configs(args, {}) + + assert isinstance(configs[0], SVDQuantConfig) + assert configs[0].rank == 32 + assert configs[0].smooth_enabled is False + assert configs[0].residual_iters == 1 + assert configs[0].model_adapter == "auto" + assert isinstance(configs[1], SignRoundConfig) + + def _normalize_options(raw): if raw is None: return None From 0649595ccca622451a270cfb21306ce50fdf490a Mon Sep 17 00:00:00 2001 From: changwangss Date: Tue, 28 Jul 2026 16:37:50 +0800 Subject: [PATCH 4/9] feat: add SVDQuant smooth calibration Signed-off-by: changwangss --- .../algorithms/transforms/svdquant/apply.py | 425 +++++++++++++++++- .../algorithms/transforms/svdquant/smooth.py | 151 +++++++ .../svdquant/smooth_adapters/__init__.py | 4 +- .../svdquant/smooth_adapters/base.py | 47 +- .../svdquant/smooth_adapters/flux.py | 2 + .../algorithms/test_svdquant_smooth.py | 156 +++++++ .../test_svdquant_smooth_adapters.py | 74 +++ 7 files changed, 850 insertions(+), 9 deletions(-) create mode 100644 auto_round/algorithms/transforms/svdquant/smooth.py create mode 100644 test/test_cpu/algorithms/test_svdquant_smooth.py create mode 100644 test/test_cpu/algorithms/test_svdquant_smooth_adapters.py diff --git a/auto_round/algorithms/transforms/svdquant/apply.py b/auto_round/algorithms/transforms/svdquant/apply.py index 79116e2dd9..f24086665e 100644 --- a/auto_round/algorithms/transforms/svdquant/apply.py +++ b/auto_round/algorithms/transforms/svdquant/apply.py @@ -14,16 +14,33 @@ from __future__ import annotations +import math +import random +from dataclasses import dataclass, field +from functools import partial +from typing import Any + import torch +import auto_round.algorithms.transforms.svdquant.residual as residual_module from auto_round.algorithms.registry import register_pipeline_member from auto_round.algorithms.transforms.base import BasePreprocessor from auto_round.algorithms.transforms.svdquant.config import SVDQuantConfig from auto_round.algorithms.transforms.svdquant.residual import ( + ActivationQuantScheme, ResidualQuantScheme, iterate_residual_decomposition, truncated_svd, ) +from auto_round.algorithms.transforms.svdquant.smooth import ( + SmoothCandidate, + absmax_channel_span, + build_alpha_beta_candidates, + build_smooth_scale, + select_best_layer_candidate, + summarize_smooth_scale, + validate_smooth_scale_for_deployment, +) from auto_round.algorithms.transforms.svdquant.smooth_adapters import SmoothSearchGroup, discover_svdquant_groups from auto_round.algorithms.transforms.svdquant.wrapper import SVDQuantLinear from auto_round.logger import logger @@ -48,6 +65,75 @@ _MXFP4_ALIASES = frozenset({"mx_fp", "mx_fp4", "mx_fp4e2m1"}) +def _detach_to_cpu(value: Any) -> Any: + if torch.is_tensor(value): + return value.detach().to("cpu", copy=True) + if isinstance(value, tuple): + return tuple(_detach_to_cpu(item) for item in value) + if isinstance(value, list): + return [_detach_to_cpu(item) for item in value] + if isinstance(value, dict): + return {key: _detach_to_cpu(item) for key, item in value.items()} + return value + + +def _move_to_device(value: Any, device: torch.device, dtype: torch.dtype | None = None) -> Any: + if torch.is_tensor(value): + target_dtype = dtype if dtype is not None and value.is_floating_point() else value.dtype + return value.to(device=device, dtype=target_dtype) + if isinstance(value, tuple): + return tuple(_move_to_device(item, device, dtype) for item in value) + if isinstance(value, list): + return [_move_to_device(item, device, dtype) for item in value] + if isinstance(value, dict): + return {key: _move_to_device(item, device, dtype) for key, item in value.items()} + return value + + +@dataclass +class CapturedEvaluation: + args: tuple[Any, ...] + kwargs: dict[str, Any] + output: Any + + +@dataclass +class SmoothGroupCalibration: + group: SmoothSearchGroup + limit: int + projection_inputs: list[torch.Tensor] = field(default_factory=list) + evaluation_calls: list[CapturedEvaluation] = field(default_factory=list) + seen_calls: int = 0 + pending_slot: int | None = None + pending_input: torch.Tensor | None = None + random: random.Random = field(default_factory=lambda: random.Random(0)) + + def begin_call(self, inputs: torch.Tensor) -> None: + self.seen_calls += 1 + if len(self.projection_inputs) < self.limit: + slot = len(self.projection_inputs) + else: + candidate = self.random.randrange(self.seen_calls) + slot = candidate if candidate < self.limit else None + self.pending_slot = slot + self.pending_input = _detach_to_cpu(inputs) if slot is not None else None + + def finish_call(self, args: tuple[Any, ...], kwargs: dict[str, Any], output: Any) -> None: + slot = self.pending_slot + captured_input = self.pending_input + self.pending_slot = None + self.pending_input = None + if slot is None or captured_input is None: + return + captured = CapturedEvaluation(_detach_to_cpu(args), _detach_to_cpu(kwargs), _detach_to_cpu(output)) + if slot == len(self.projection_inputs): + self.projection_inputs.append(captured_input) + self.evaluation_calls.append(captured) + else: + self.projection_inputs[slot] = captured_input + self.evaluation_calls[slot] = captured + + @register_pipeline_member(SVDQuantConfig) class SVDQuantTransform(BasePreprocessor): """Split target linears into a quantized residual and an FP low-rank branch.""" @@ -56,6 +142,7 @@ def __init__(self, config: SVDQuantConfig) -> None: super().__init__(config) self._configured_block_names: tuple[str, ...] = () self._block_groups: dict[str, list[SmoothSearchGroup]] = {} + self._smooth_calibration: dict[str, SmoothGroupCalibration] = {} def bind(self, orchestrator) -> None: super().bind(orchestrator) @@ -81,17 +168,44 @@ def prepare_run(self, composer=None) -> None: ) def register_fp_input_forward_hooks(self, block) -> list: - if self.config.smooth_enabled: - raise NotImplementedError("SVDQuant smooth calibration is not ported to the main architecture yet.") - return [] + if not self.config.smooth_enabled: + return [] + self._clear_smooth_calibration() + block_name = str(getattr(block, "global_name", "")) + groups = self._block_groups.get(block_name) + if groups is None: + groups = discover_svdquant_groups(block, self._is_target) + self._block_groups[block_name] = groups + self._smooth_calibration = { + group.key: SmoothGroupCalibration(group, self.config.smooth_max_calibration_calls) for group in groups + } + + projection_owners = {} + evaluation_owners = {} + modules = {} + for calibration in self._smooth_calibration.values(): + projection_module = calibration.group.projection_input_module + evaluation_module = calibration.group.evaluation_module + projection_owners.setdefault(id(projection_module), []).append(calibration) + evaluation_owners.setdefault(id(evaluation_module), []).append(calibration) + modules[id(projection_module)] = projection_module + modules[id(evaluation_module)] = evaluation_module + + def collect_calibration(module, inputs, kwargs, output): + for calibration in projection_owners.get(id(module), ()): + if inputs and torch.is_tensor(inputs[0]): + value = inputs[0] + if value.shape[-1] == calibration.group.projections[0].in_features: + calibration.begin_call(value) + for calibration in evaluation_owners.get(id(module), ()): + calibration.finish_call(inputs, kwargs, output) + + return [module.register_forward_hook(collect_calibration, with_kwargs=True) for module in modules.values()] @torch.no_grad() def pre_quantize_block(self, ctx) -> None: if len(ctx.block_names) != 1: raise ValueError(f"SVDQuant requires one block at a time, got {ctx.block_names!r}.") - if self.config.smooth_enabled: - raise NotImplementedError("SVDQuant smooth calibration is not ported to the main architecture yet.") - block_name = ctx.block_name block = ctx.model.get_submodule(block_name) groups = self._block_groups.get(block_name) @@ -99,6 +213,10 @@ def pre_quantize_block(self, ctx) -> None: groups = discover_svdquant_groups(block, self._is_target) self._block_groups[block_name] = groups + if self.config.smooth_enabled: + self._pre_quantize_smoothed_block(block, groups) + return + local_names = {id(module): name for name, module in block.named_modules() if name} replacements = [] for group in groups: @@ -113,11 +231,266 @@ def pre_quantize_block(self, ctx) -> None: _set_child_module(block, local_name, wrapper) def post_quantize_block(self, ctx) -> None: + self._clear_smooth_calibration() self._block_groups.pop(ctx.block_name, None) def finalize_run(self) -> None: + self._clear_smooth_calibration() self._block_groups.clear() + def _clear_smooth_calibration(self) -> None: + self._smooth_calibration.clear() + + def _pre_quantize_smoothed_block(self, block: torch.nn.Module, groups: list[SmoothSearchGroup]) -> None: + if not self._smooth_calibration: + raise ValueError("SVDQuant smooth calibration inputs are missing for the current block.") + local_names = {id(module): name for name, module in block.named_modules() if name} + try: + selected_scales = {} + for group in groups: + calibration = self._smooth_calibration.get(group.key) + if calibration is None or not calibration.projection_inputs or not calibration.evaluation_calls: + raise ValueError(f"SVDQuant smooth calibration inputs are missing for group {group.key!r}.") + selected_scales[group.key] = self._search_group_scale(calibration, block, local_names) + + replacements = [] + for group in groups: + calibration = self._smooth_calibration[group.key] + wrappers = self._decompose_smoothed_group(calibration, selected_scales[group.key], block, local_names) + for projection, wrapper in zip(group.projections, wrappers): + local_name = local_names.get(id(projection)) + if local_name is None: + raise ValueError(f"SVDQuant could not locate projection {self._module_name(projection)!r}.") + replacements.append((local_name, wrapper)) + for local_name, wrapper in replacements: + _set_child_module(block, local_name, wrapper) + finally: + self._clear_smooth_calibration() + + def _search_group_scale( + self, + calibration: SmoothGroupCalibration, + block: torch.nn.Module, + local_names: dict[int, str], + ) -> torch.Tensor: + group = calibration.group + device = group.projections[0].weight.device + x_span = torch.stack([absmax_channel_span(inputs, -1) for inputs in calibration.projection_inputs], dim=0).amax( + dim=0 + ) + weights = [ + projection.weight.detach().to(device=device, dtype=torch.float32) for projection in group.projections + ] + w_span = absmax_channel_span(torch.cat(weights, dim=0), 1).cpu() + scored = [] + for alpha, beta in build_alpha_beta_candidates(self.config.smooth_num_grids): + scale = build_smooth_scale(x_span, w_span, alpha, beta, eps=self.config.smooth_eps) + try: + scale = validate_smooth_scale_for_deployment( + scale, dtype=group.projections[0].weight.dtype, module_name=group.key + ).to(torch.float32) + error = self._score_group_wrappers( + calibration, self._candidate_group_wrappers(group, scale), block, local_names + ) + except (RuntimeError, ValueError, TypeError) as exc: + logger.debug("Skipping SVDQuant smooth candidate (%s, %s) for %s: %s", alpha, beta, group.key, exc) + error = float("inf") + candidate = SmoothCandidate(alpha, beta, scale) + scored.append(((candidate, error), error)) + selected, error = select_best_layer_candidate(scored, module_name=group.key) + self._log_selected_smooth_candidate(group.key, selected, error) + return selected.scale + + @staticmethod + def _log_selected_smooth_candidate(module_name: str, candidate: SmoothCandidate, error: float) -> None: + stats = summarize_smooth_scale(candidate.scale) + logger.info( + "SVDQuant smooth selected for %s: alpha=%.6g beta=%.6g error=%.6g " + "scale_min=%.6g scale_max=%.6g scale_ratio=%.6g below_1e-3=%d above_20=%d", + module_name, + candidate.alpha, + candidate.beta, + error, + stats.minimum, + stats.maximum, + stats.ratio, + stats.below_min_count, + stats.above_max_count, + ) + if stats.below_min_count or stats.above_max_count: + logger.warning( + "SVDQuant smooth scale for %s contains extreme values: below_1e-3=%d above_20=%d", + module_name, + stats.below_min_count, + stats.above_max_count, + ) + + def _score_group_wrappers( + self, + calibration: SmoothGroupCalibration, + wrappers: list[SVDQuantLinear], + block: torch.nn.Module, + local_names: dict[int, str], + ) -> float: + group = calibration.group + replacements = [] + for projection, wrapper in zip(group.projections, wrappers): + local_name = local_names.get(id(projection)) + if local_name is None: + raise ValueError(f"SVDQuant could not locate projection {self._module_name(projection)!r}.") + replacements.append((local_name, projection, wrapper)) + try: + for local_name, _, wrapper in replacements: + _set_child_module(block, local_name, wrapper) + error = torch.zeros((), dtype=torch.float64) + for call in calibration.evaluation_calls: + evaluation_module = group.evaluation_module + if len(group.projections) == 1 and evaluation_module is group.projections[0]: + evaluation_module = wrappers[0] + device = group.projections[0].weight.device + dtype = group.projections[0].weight.dtype + args = _move_to_device(call.args, device, dtype) + kwargs = group.filter_evaluation_kwargs(_move_to_device(call.kwargs, device, dtype)) + actual = group.normalize_output(evaluation_module(*args, **kwargs)) + reference = tuple(tensor.to(device) for tensor in group.normalize_output(call.output)) + if len(actual) != len(reference): + raise ValueError("SVDQuant smooth output tensor count changed.") + for actual_tensor, reference_tensor in zip(actual, reference): + if actual_tensor.shape != reference_tensor.shape: + raise ValueError("SVDQuant smooth output tensor shape changed.") + error += torch.sum((actual_tensor.float() - reference_tensor.float()).square()).double().cpu() + return error.item() + finally: + for local_name, projection, _ in replacements: + _set_child_module(block, local_name, projection) + + def _candidate_group_wrappers(self, group: SmoothSearchGroup, scale: torch.Tensor) -> list[SVDQuantLinear]: + weights = [ + projection.weight.detach().to(torch.float32) * scale.to(projection.weight.device) + for projection in group.projections + ] + stacked = torch.cat(weights, dim=0) + output_sizes = [projection.out_features for projection in group.projections] + rank = min(self.config.rank, *stacked.shape) + _, down, up = truncated_svd(stacked, rank) + low_rank_dtype = self._resolve_low_rank_dtype(group.projections[0].weight.dtype) + deployed_down = down.to(low_rank_dtype) + deployed_up = up.to(low_rank_dtype) + low_rank_parts = (deployed_up.float() @ deployed_down.float()).split(output_sizes, dim=0) + residuals = [] + for projection, weight, low_rank in zip(group.projections, weights, low_rank_parts): + residual = (weight - low_rank).to(projection.weight.dtype) + residuals.append(residual_module.rtn_qdq_residual(residual, self._residual_quant_scheme(projection))) + return self._build_group_wrappers( + group, + residuals, + deployed_down, + deployed_up.split(output_sizes, dim=0), + scale, + activation_scheme=self._group_activation_quant_scheme(group), + ) + + def _decompose_smoothed_group( + self, + calibration: SmoothGroupCalibration, + scale: torch.Tensor, + block: torch.nn.Module, + local_names: dict[int, str], + ) -> list[SVDQuantLinear]: + group = calibration.group + weights = [ + projection.weight.detach().to(torch.float32) * scale.to(projection.weight.device) + for projection in group.projections + ] + stacked = torch.cat(weights, dim=0) + output_sizes = [projection.out_features for projection in group.projections] + rank = min(self.config.rank, *stacked.shape) + low_rank_dtype = self._resolve_low_rank_dtype(group.projections[0].weight.dtype) + if self.config.residual_iters == 1: + _, down, up = truncated_svd(stacked, rank) + deployed_down = down.to(low_rank_dtype) + deployed_up = up.to(low_rank_dtype) + low_rank = deployed_up.float() @ deployed_down.float() + residuals = [ + residual.to(projection.weight.dtype) + for residual, projection in zip(stacked.sub(low_rank).split(output_sizes, dim=0), group.projections) + ] + else: + residuals, deployed_down, deployed_up = self._iterate_smoothed_group_residual( + calibration, block, local_names, stacked, rank, low_rank_dtype, scale + ) + return self._build_group_wrappers( + group, residuals, deployed_down, deployed_up.split(output_sizes, dim=0), scale + ) + + def _iterate_smoothed_group_residual( + self, + calibration: SmoothGroupCalibration, + block: torch.nn.Module, + local_names: dict[int, str], + stacked: torch.Tensor, + rank: int, + low_rank_dtype: torch.dtype, + scale: torch.Tensor, + ) -> tuple[list[torch.Tensor], torch.Tensor, torch.Tensor]: + group = calibration.group + output_sizes = [projection.out_features for projection in group.projections] + quantized_residual = torch.zeros_like(stacked) + best = None + best_error = float("inf") + activation_scheme = self._group_activation_quant_scheme(group) + for iteration in range(1, self.config.residual_iters + 1): + _, down, up = truncated_svd(stacked - quantized_residual, rank) + deployed_down = down.to(low_rank_dtype) + deployed_up = up.to(low_rank_dtype) + low_rank = deployed_up.float() @ deployed_down.float() + residual_parts = (stacked - low_rank).split(output_sizes, dim=0) + qdq_residuals = [ + residual_module.rtn_qdq_residual( + residual.to(projection.weight.dtype), self._residual_quant_scheme(projection) + ) + for residual, projection in zip(residual_parts, group.projections) + ] + quantized_residual = torch.cat([residual.float() for residual in qdq_residuals], dim=0) + wrappers = self._build_group_wrappers( + group, + qdq_residuals, + deployed_down, + deployed_up.split(output_sizes, dim=0), + scale, + activation_scheme=activation_scheme, + ) + error = self._score_group_wrappers(calibration, wrappers, block, local_names) + accepted = math.isfinite(error) and error <= best_error + if accepted: + best = (deployed_down.clone(), deployed_up.clone(), iteration) + best_error = error + elif self.config.residual_early_stop and best is not None: + logger.info( + "SVDQuant residual early stop for %s at iteration %d: output error %.6g > best %.6g", + group.key, + iteration, + error, + best_error, + ) + break + if best is None: + raise ValueError(f"SVDQuant residual iteration failed for group {group.key!r}.") + deployed_down, deployed_up, selected_iteration = best + logger.info( + "SVDQuant residual selected iteration %d/%d for %s: output error %.6g", + selected_iteration, + self.config.residual_iters, + group.key, + best_error, + ) + low_rank = deployed_up.float() @ deployed_down.float() + residuals = [ + residual.to(projection.weight.dtype) + for residual, projection in zip((stacked - low_rank).split(output_sizes, dim=0), group.projections) + ] + return residuals, deployed_down, deployed_up + def _decompose_group(self, group: SmoothSearchGroup) -> list[SVDQuantLinear]: weights = [projection.weight.detach().to(torch.float32) for projection in group.projections] stacked = torch.cat(weights, dim=0) @@ -172,9 +545,11 @@ def _build_group_wrappers( down: torch.Tensor, up_parts: tuple[torch.Tensor, ...], smooth: torch.Tensor, + activation_scheme: ActivationQuantScheme | None = None, ) -> list[SVDQuantLinear]: wrappers = [] rank = down.shape[0] + input_smooth = smooth.reciprocal() for projection, residual_weight, up in zip(group.projections, residuals, up_parts): residual = self._new_linear_like(projection, residual_weight, projection.bias) lora_down = torch.nn.Linear( @@ -196,7 +571,20 @@ def _build_group_wrappers( self._mark_unquantized(lora_down) self._mark_unquantized(lora_up) self._copy_quant_attrs(projection, residual, suffix=".residual_linear") - wrappers.append(SVDQuantLinear(residual, lora_down, lora_up, smooth.to(projection.weight.dtype))) + activation_qdq = ( + None + if activation_scheme is None + else partial(residual_module.rtn_qdq_activation, scheme=activation_scheme) + ) + wrappers.append( + SVDQuantLinear( + residual, + lora_down, + lora_up, + input_smooth.to(projection.weight.dtype), + activation_qdq=activation_qdq, + ) + ) return wrappers def _shared_residual_scheme(self, group: SmoothSearchGroup) -> ResidualQuantScheme: @@ -205,6 +593,29 @@ def _shared_residual_scheme(self, group: SmoothSearchGroup) -> ResidualQuantSche raise ValueError(f"SVDQuant group {group.key!r} has inconsistent residual quantization schemes.") return schemes[0] + def _group_activation_quant_scheme(self, group: SmoothSearchGroup) -> ActivationQuantScheme: + schemes = tuple(self._activation_quant_scheme(projection) for projection in group.projections) + if any(scheme != schemes[0] for scheme in schemes[1:]): + raise ValueError(f"SVDQuant group {group.key!r} has inconsistent activation quantization schemes.") + return schemes[0] + + def _activation_quant_scheme(self, module: torch.nn.Linear) -> ActivationQuantScheme: + attributes = { + "data_type": "act_data_type", + "bits": "act_bits", + "group_size": "act_group_size", + "sym": "act_sym", + } + missing = [ + source for source in attributes.values() if not hasattr(module, source) or getattr(module, source) is None + ] + if missing: + raise ValueError( + f"SVDQuant smooth calibration requires a complete activation scheme for " + f"{self._module_name(module)!r}; missing: {', '.join(missing)}." + ) + return ActivationQuantScheme(**{target: getattr(module, source) for target, source in attributes.items()}) + def _residual_quant_scheme(self, module: torch.nn.Linear) -> ResidualQuantScheme: required = ("data_type", "bits", "group_size", "sym") missing = [attr for attr in required if not hasattr(module, attr) or getattr(module, attr) is None] diff --git a/auto_round/algorithms/transforms/svdquant/smooth.py b/auto_round/algorithms/transforms/svdquant/smooth.py new file mode 100644 index 0000000000..4c37b06b7b --- /dev/null +++ b/auto_round/algorithms/transforms/svdquant/smooth.py @@ -0,0 +1,151 @@ +# Copyright (c) 2026 Intel Corporation +# +# 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. + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Iterable, TypeVar + +import torch + + +@dataclass(frozen=True) +class SmoothCandidate: + alpha: float + beta: float + scale: torch.Tensor + + +@dataclass(frozen=True) +class SmoothScaleStats: + minimum: float + maximum: float + ratio: float + below_min_count: int + above_max_count: int + + +def build_alpha_beta_candidates(num_grids: int) -> list[tuple[float, float]]: + if type(num_grids) is not int or num_grids < 2: + raise ValueError(f"`num_grids` must be an integer greater than or equal to 2, got {num_grids!r}") + choices = [index / num_grids for index in range(1, num_grids)] + return [(0.0, 0.0), *[(alpha, 0.0) for alpha in choices], *[(alpha, 1.0 - alpha) for alpha in choices]] + + +def absmax_channel_span(tensor: torch.Tensor, channels_dim: int) -> torch.Tensor: + if tensor.ndim == 0: + raise ValueError("Cannot calculate a channel span for a scalar tensor.") + channels_dim %= tensor.ndim + moved = tensor.detach().movedim(channels_dim, -1) + return moved.abs().reshape(-1, moved.shape[-1]).amax(dim=0).to(torch.float32) + + +def build_smooth_scale( + x_span: torch.Tensor, + w_span: torch.Tensor, + alpha: float, + beta: float, + eps: float | None = None, +) -> torch.Tensor: + if not 0.0 <= alpha <= 1.0 or not 0.0 <= beta <= 1.0: + raise ValueError(f"Smooth alpha and beta must be in [0, 1], got alpha={alpha!r}, beta={beta!r}") + if x_span.shape != w_span.shape: + raise ValueError(f"Smooth spans must have matching shapes, got {x_span.shape} and {w_span.shape}") + + x_span = x_span.to(torch.float32) + w_span = w_span.to(device=x_span.device, dtype=torch.float32) + x_zero = x_span == 0 + w_zero = w_span == 0 + if eps is not None: + if eps <= 0: + raise ValueError(f"`eps` must be positive, got {eps!r}") + x_span = torch.where(x_zero, eps, x_span) + w_span = torch.where(w_zero, eps, w_span) + if alpha == 0.0 and beta == 0.0: + return torch.ones_like(x_span) + + if alpha > 0.0: + scale = x_span.pow(alpha) + if beta > 0.0: + scale = scale / w_span.pow(beta) + else: + scale = w_span.pow(-beta) + + scale = scale.clone() + if beta > 0.0 and bool(w_zero.any()): + scale.fill_(1) + elif alpha > 0.0: + scale[x_zero] = 1 + scale[scale == 0] = 1 + if not torch.isfinite(scale).all(): + scale.fill_(1) + return scale + + +def validate_smooth_scale_for_deployment( + scale: torch.Tensor, + *, + dtype: torch.dtype, + module_name: str, +) -> torch.Tensor: + """Validate a smooth scale after materialization in its deployment dtype.""" + deployed = scale.to(dtype=dtype) + reciprocal = deployed.reciprocal() + if ( + not bool(torch.isfinite(deployed).all()) + or not bool((deployed > 0).all()) + or not bool(torch.isfinite(reciprocal).all()) + or not bool((reciprocal > 0).all()) + ): + raise ValueError(f"SVDQuant smooth scale is not deployable for {module_name!r} in dtype {dtype}.") + return deployed + + +def summarize_smooth_scale( + scale: torch.Tensor, + *, + low_threshold: float = 1e-3, + high_threshold: float = 20.0, +) -> SmoothScaleStats: + values = scale.detach().to(device="cpu", dtype=torch.float32) + minimum = values.amin().item() + maximum = values.amax().item() + return SmoothScaleStats( + minimum=minimum, + maximum=maximum, + ratio=maximum / minimum, + below_min_count=int((values < low_threshold).sum().item()), + above_max_count=int((values > high_threshold).sum().item()), + ) + + +_CandidateT = TypeVar("_CandidateT") + + +def select_best_layer_candidate( + candidates: Iterable[tuple[_CandidateT, float | torch.Tensor]], + *, + module_name: str, +) -> _CandidateT: + best_candidate = None + best_error = float("inf") + for candidate, error in candidates: + error_value = error.item() if torch.is_tensor(error) else float(error) + if math.isfinite(error_value) and error_value <= best_error: + best_candidate = candidate + best_error = error_value + if best_candidate is None: + raise ValueError(f"SVDQuant smooth search produced no finite candidate for {module_name!r}.") + return best_candidate diff --git a/auto_round/algorithms/transforms/svdquant/smooth_adapters/__init__.py b/auto_round/algorithms/transforms/svdquant/smooth_adapters/__init__.py index b79f0b0e41..c8613c3052 100644 --- a/auto_round/algorithms/transforms/svdquant/smooth_adapters/__init__.py +++ b/auto_round/algorithms/transforms/svdquant/smooth_adapters/__init__.py @@ -31,4 +31,6 @@ def discover_svdquant_groups(block: torch.nn.Module, is_target: TargetPredicate) return generic_linear_groups(block, is_target) -__all__ = ["SmoothSearchGroup", "discover_svdquant_groups"] +discover_smooth_search_groups = discover_svdquant_groups + +__all__ = ["SmoothSearchGroup", "discover_smooth_search_groups", "discover_svdquant_groups"] diff --git a/auto_round/algorithms/transforms/svdquant/smooth_adapters/base.py b/auto_round/algorithms/transforms/svdquant/smooth_adapters/base.py index 34a9179407..30f01bfa3c 100644 --- a/auto_round/algorithms/transforms/svdquant/smooth_adapters/base.py +++ b/auto_round/algorithms/transforms/svdquant/smooth_adapters/base.py @@ -14,14 +14,49 @@ from __future__ import annotations -from collections.abc import Callable +import inspect +from collections.abc import Callable, Mapping from dataclasses import dataclass +from typing import Any import torch TargetPredicate = Callable[[str, torch.nn.Module], bool] +def normalize_tensors(output: Any, indices: tuple[int, ...] | None = None) -> tuple[torch.Tensor, ...]: + """Return floating-point tensors used by the smooth output-error objective.""" + if indices is not None: + if not isinstance(output, (tuple, list)): + raise TypeError("Indexed SVDQuant smooth output must be a tuple or list.") + output = tuple(output[index] for index in indices) + + tensors = [] + + def collect(value): + if torch.is_tensor(value): + if value.is_floating_point(): + tensors.append(value) + elif isinstance(value, Mapping): + for item in value.values(): + collect(item) + elif isinstance(value, (tuple, list)): + for item in value: + collect(item) + + collect(output) + if not tensors: + raise ValueError("SVDQuant smooth evaluation produced no floating-point tensors.") + return tuple(tensors) + + +def filter_supported_kwargs(module: torch.nn.Module, kwargs: Mapping[str, Any]) -> dict[str, Any]: + parameters = inspect.signature(module.forward).parameters + if any(parameter.kind is inspect.Parameter.VAR_KEYWORD for parameter in parameters.values()): + return dict(kwargs) + return {name: value for name, value in kwargs.items() if name in parameters} + + @dataclass(frozen=True) class SmoothSearchGroup: """Projection group that shares an input scale and low-rank down factor.""" @@ -29,7 +64,9 @@ class SmoothSearchGroup: key: str projection_names: tuple[str, ...] projections: tuple[torch.nn.Linear, ...] + projection_input_key: str projection_input_module: torch.nn.Module + evaluation_input_key: str evaluation_module: torch.nn.Module output_indices: tuple[int, ...] | None = None output_splits: tuple[int, ...] = () @@ -44,6 +81,12 @@ def __post_init__(self) -> None: if self.output_splits and sum(self.output_splits) != len(self.projections): raise ValueError(f"SVDQuant group {self.key!r} output splits do not cover its projections.") + def filter_evaluation_kwargs(self, kwargs: Mapping[str, Any]) -> dict[str, Any]: + return filter_supported_kwargs(self.evaluation_module, kwargs) + + def normalize_output(self, output: Any) -> tuple[torch.Tensor, ...]: + return normalize_tensors(output, self.output_indices) + def module_global_name(block: torch.nn.Module, local_name: str) -> str: prefix = str(getattr(block, "global_name", block.__class__.__name__)) @@ -85,7 +128,9 @@ def generic_linear_groups( key=name, projection_names=(name,), projections=(module,), + projection_input_key=name, projection_input_module=module, + evaluation_input_key=name, evaluation_module=module, output_splits=(1,), ) diff --git a/auto_round/algorithms/transforms/svdquant/smooth_adapters/flux.py b/auto_round/algorithms/transforms/svdquant/smooth_adapters/flux.py index e7b4d5d30f..be680e078d 100644 --- a/auto_round/algorithms/transforms/svdquant/smooth_adapters/flux.py +++ b/auto_round/algorithms/transforms/svdquant/smooth_adapters/flux.py @@ -67,7 +67,9 @@ def _make_group( key=module_global_name(block, key), projection_names=names, projections=projections, + projection_input_key=names[0], projection_input_module=projections[0], + evaluation_input_key=module_global_name(block, evaluation_path), evaluation_module=evaluation_module, output_indices=output_indices, output_splits=splits, diff --git a/test/test_cpu/algorithms/test_svdquant_smooth.py b/test/test_cpu/algorithms/test_svdquant_smooth.py new file mode 100644 index 0000000000..d756b54f70 --- /dev/null +++ b/test/test_cpu/algorithms/test_svdquant_smooth.py @@ -0,0 +1,156 @@ +# Copyright (c) 2026 Intel Corporation +# +# 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. + +from types import SimpleNamespace + +import pytest +import torch + +import auto_round.algorithms.transforms.svdquant.residual as residual_module +from auto_round.algorithms.composer import BlockContext +from auto_round.algorithms.transforms.svdquant import SVDQuantConfig, SVDQuantLinear, SVDQuantTransform +from auto_round.algorithms.transforms.svdquant.smooth import ( + absmax_channel_span, + build_alpha_beta_candidates, + build_smooth_scale, + select_best_layer_candidate, + validate_smooth_scale_for_deployment, +) + + +def test_alpha_beta_grid_matches_proven_candidate_order(): + assert build_alpha_beta_candidates(4) == [ + (0.0, 0.0), + (0.25, 0.0), + (0.5, 0.0), + (0.75, 0.0), + (0.25, 0.75), + (0.5, 0.5), + (0.75, 0.25), + ] + assert len(build_alpha_beta_candidates(20)) == 39 + + +def test_smooth_scale_uses_activation_and_weight_channel_spans(): + activations = torch.tensor([[[1.0, -9.0, 4.0], [16.0, 2.0, -1.0]]]) + weights = torch.tensor([[1.0, 4.0, 16.0], [-0.5, 2.0, 8.0]]) + x_span = absmax_channel_span(activations, -1) + w_span = absmax_channel_span(weights, 1) + + scale = build_smooth_scale(x_span, w_span, alpha=0.5, beta=0.5) + + torch.testing.assert_close(scale, x_span.sqrt() / w_span.sqrt()) + + +def test_smooth_scale_zero_channels_follow_identity_fallback(): + scale = build_smooth_scale( + torch.tensor([0.0, 4.0]), + torch.tensor([0.0, 1.0]), + alpha=0.5, + beta=0.5, + eps=1e-6, + ) + + torch.testing.assert_close(scale, torch.ones(2)) + + +def test_deployment_validation_rejects_bfloat16_reciprocal_overflow(): + with pytest.raises(ValueError, match="deployable.*proj"): + validate_smooth_scale_for_deployment( + torch.tensor([1e-40, 1.0]), + dtype=torch.bfloat16, + module_name="transformer_blocks.0.proj", + ) + + +def test_candidate_selection_keeps_later_exact_tie_and_skips_nonfinite(): + candidates = [("nan", float("nan")), ("first", 1.0), ("later", 1.0), ("worse", 2.0)] + + assert select_best_layer_candidate(candidates, module_name="blocks.0.qkv") == "later" + + +class SmoothBlock(torch.nn.Module): + def __init__(self): + super().__init__() + self.proj = torch.nn.Linear(4, 3, bias=False) + + def forward(self, hidden_states): + return self.proj(hidden_states) + + +class SmoothModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.blocks = torch.nn.ModuleList([SmoothBlock()]) + + +def test_smooth_calibration_reservoir_is_bounded_and_preserves_forward(monkeypatch): + model = SmoothModel() + block_name = "blocks.0" + for name, module in model.named_modules(): + module.global_name = name + if isinstance(module, torch.nn.Linear): + module.data_type = "mx_fp4e2m1" + module.bits = 4 + module.group_size = 32 + module.sym = True + module.act_data_type = "mx_fp4e2m1" + module.act_bits = 4 + module.act_group_size = 32 + module.act_sym = True + + monkeypatch.setattr(residual_module, "rtn_qdq_residual", lambda tensor, _scheme: tensor) + monkeypatch.setattr(residual_module, "rtn_qdq_activation", lambda tensor, scheme: tensor) + transform = SVDQuantTransform( + SVDQuantConfig( + rank=1, + smooth_enabled=True, + smooth_num_grids=2, + smooth_max_calibration_calls=2, + low_rank_dtype="fp32", + ) + ) + transform.bind( + SimpleNamespace( + model_context=SimpleNamespace(model=model), + compress_context=None, + calibration_context=None, + scheme_context=None, + scale_dtype=None, + nblocks=1, + quant_block_list=[[block_name]], + ) + ) + transform.prepare_run() + block = model.blocks[0] + handles = transform.register_fp_input_forward_hooks(block) + for index in range(10): + block(torch.full((1, 4), float(index + 1))) + for handle in handles: + handle.remove() + + calibration = next(iter(transform._smooth_calibration.values())) + assert calibration.seen_calls == 10 + assert len(calibration.projection_inputs) == 2 + assert len(calibration.evaluation_calls) == 2 + inputs = torch.randn(2, 4) + expected = block(inputs) + + transform.pre_quantize_block( + BlockContext(model=model, block_names=[block_name], block_name=block_name, block_index=0) + ) + + assert isinstance(block.proj, SVDQuantLinear) + torch.testing.assert_close(block(inputs), expected, atol=1e-5, rtol=1e-5) + assert not transform._smooth_calibration diff --git a/test/test_cpu/algorithms/test_svdquant_smooth_adapters.py b/test/test_cpu/algorithms/test_svdquant_smooth_adapters.py new file mode 100644 index 0000000000..6e798a2748 --- /dev/null +++ b/test/test_cpu/algorithms/test_svdquant_smooth_adapters.py @@ -0,0 +1,74 @@ +# Copyright (c) 2026 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +from auto_round.algorithms.transforms.svdquant.smooth_adapters import discover_svdquant_groups + + +class Attention(torch.nn.Module): + def __init__(self, width): + super().__init__() + self.to_q = torch.nn.Linear(width, width) + self.to_k = torch.nn.Linear(width, width) + self.to_v = torch.nn.Linear(width, width) + self.add_q_proj = torch.nn.Linear(width, width) + self.add_k_proj = torch.nn.Linear(width, width) + self.add_v_proj = torch.nn.Linear(width, width) + self.to_out = torch.nn.ModuleList([torch.nn.Linear(width, width)]) + self.to_add_out = torch.nn.Linear(width, width) + + +class Projection(torch.nn.Module): + def __init__(self, in_features, out_features): + super().__init__() + self.proj = torch.nn.Linear(in_features, out_features) + + +class FeedForward(torch.nn.Module): + def __init__(self, width): + super().__init__() + self.net = torch.nn.ModuleList( + [Projection(width, width * 2), torch.nn.GELU(), torch.nn.Linear(width * 2, width)] + ) + + +class FluxTransformerBlock(torch.nn.Module): + def __init__(self, width=4): + super().__init__() + self.attn = Attention(width) + self.ff = FeedForward(width) + self.ff_context = FeedForward(width) + + +def test_flux_adapter_discovers_fused_projection_groups_without_duplicates(): + block = FluxTransformerBlock() + block.global_name = "transformer_blocks.7" + + groups = discover_svdquant_groups(block, lambda _name, module: isinstance(module, torch.nn.Linear)) + by_key = {group.key: group for group in groups} + + assert set(by_key) == { + "transformer_blocks.7.attn.qkv", + "transformer_blocks.7.attn.add_qkv", + "transformer_blocks.7.attn.to_out.0", + "transformer_blocks.7.attn.to_add_out", + "transformer_blocks.7.ff.net.0.proj", + "transformer_blocks.7.ff.net.2", + "transformer_blocks.7.ff_context.net.0.proj", + "transformer_blocks.7.ff_context.net.2", + } + assert len(by_key["transformer_blocks.7.attn.qkv"].projections) == 3 + assert len(by_key["transformer_blocks.7.attn.add_qkv"].projections) == 3 + assert len({id(projection) for group in groups for projection in group.projections}) == 12 From bc1997baba94ca5d2d94cb15d316b982c4b90b86 Mon Sep 17 00:00:00 2001 From: changwangss Date: Tue, 28 Jul 2026 16:40:58 +0800 Subject: [PATCH 5/9] feat: add SVDQuant serialization codecs Signed-off-by: changwangss --- auto_round/export/svdquant_mxfp4.py | 404 ++++++++++++++++++++ auto_round/export/svdquant_w4a16.py | 314 +++++++++++++++ test/test_cpu/export/test_svdquant_mxfp4.py | 55 +++ test/test_cpu/export/test_svdquant_w4a16.py | 56 +++ 4 files changed, 829 insertions(+) create mode 100644 auto_round/export/svdquant_mxfp4.py create mode 100644 auto_round/export/svdquant_w4a16.py create mode 100644 test/test_cpu/export/test_svdquant_mxfp4.py create mode 100644 test/test_cpu/export/test_svdquant_w4a16.py diff --git a/auto_round/export/svdquant_mxfp4.py b/auto_round/export/svdquant_mxfp4.py new file mode 100644 index 0000000000..534fdc110b --- /dev/null +++ b/auto_round/export/svdquant_mxfp4.py @@ -0,0 +1,404 @@ +# Copyright (c) 2026 Intel Corporation +# +# 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. + +"""Logical MXFP4 codecs and Nunchaku-compatible residual packing. + +E2M1 scales use an explicit rank contract: scalars are global, scales with one +fewer dimension than values are group-aligned along a new trailing singleton, +and scales with the same rank use ordinary PyTorch broadcasting. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from auto_round.data_type.mxfp import quant_element, quant_mx_rceil + +_E2M1_MAGNITUDES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) +_SUPPORTED_DECODE_DTYPES = (torch.float16, torch.bfloat16, torch.float32, torch.float64) +_SUPPORTED_PACKED_DTYPES = (torch.float16, torch.bfloat16) + + +def _validate_lowrank_weight(weight: torch.Tensor, down: bool) -> None: + if not isinstance(weight, torch.Tensor) or weight.ndim != 2: + raise ValueError("weight must be a 2D torch.Tensor") + if weight.dtype not in _SUPPORTED_PACKED_DTYPES: + raise ValueError("weight dtype must be torch.float16 or torch.bfloat16") + if not isinstance(down, bool): + raise ValueError("down must be a bool") + if 0 in weight.shape: + raise ValueError("weight dimensions must be non-empty") + if not bool(torch.isfinite(weight).all()): + raise ValueError("weight must contain only finite values") + + +def pack_lowrank_weight(weight: torch.Tensor, down: bool) -> torch.Tensor: + """Pack a logical low-rank matrix with 128-feature and 16-rank alignment.""" + + _validate_lowrank_weight(weight, down) + rows = NunchakuMXFP4Packer._ceil_to(weight.shape[0], 16 if down else 128) + columns = NunchakuMXFP4Packer._ceil_to(weight.shape[1], 128 if down else 16) + padded = torch.zeros((rows, columns), dtype=weight.dtype, device=weight.device) + padded[: weight.shape[0], : weight.shape[1]] = weight + pack_n = pack_k = 16 + if down: + rank, channels = padded.shape + rank_packs, channel_packs = rank // pack_n, channels // pack_k + packed = padded.view(rank_packs, pack_n, channel_packs, pack_k).permute(2, 0, 1, 3) + else: + channels, rank = padded.shape + channel_packs, rank_packs = channels // pack_n, rank // pack_k + packed = padded.view(channel_packs, pack_n, rank_packs, pack_k).permute(0, 2, 1, 3) + packed = packed.reshape(channel_packs, rank_packs, 2, 8, 1, 2, 4, 2) + return packed.permute(0, 1, 3, 6, 2, 5, 4, 7).contiguous().view(channels, rank) + + +def unpack_lowrank_weight(weight: torch.Tensor, down: bool) -> torch.Tensor: + """Invert :func:`pack_lowrank_weight`, retaining its padded logical shape.""" + + _validate_lowrank_weight(weight, down) + channels, rank = weight.shape + if channels % 128 or rank % 16: + raise ValueError("packed weight feature and rank dimensions must be divisible by 128 and 16 respectively") + if down: + rank_packs, channel_packs = rank // 16, channels // 16 + else: + channel_packs, rank_packs = channels // 16, rank // 16 + unpacked = weight.view(channel_packs, rank_packs, 8, 4, 2, 2, 1, 2) + unpacked = unpacked.permute(0, 1, 4, 2, 6, 5, 3, 7).contiguous().view(channel_packs, rank_packs, 16, 16) + if down: + return unpacked.permute(1, 2, 0, 3).contiguous().view(rank, channels) + return unpacked.permute(0, 2, 1, 3).contiguous().view(channels, rank) + + +@dataclass(frozen=True) +class PackedMXFP4: + """Physical MXFP4 residual tensors and their logical/padded dimensions.""" + + qweight: torch.Tensor + wscales: torch.Tensor + logical_shape: tuple[int, int] + padded_shape: tuple[int, int] + + +class NunchakuMXFP4Packer: + """Pack MXFP4 residuals into the Nunchaku 4-bit MMA memory layout. + + The reshape/permutation constants follow the Nunchaku weight and MXFP4 + micro-scale packers at reference source commit 0abaaf0. + """ + + comp_n = 16 + comp_k = mem_k = 64 + num_n_lanes = 8 + num_k_lanes = 4 + n_pack_size = 2 + k_pack_size = 2 + reg_n = 1 + reg_k = 8 + num_k_unrolls = 2 + + def __init__(self, warp_n: int = 128) -> None: + if warp_n != 128: + raise ValueError("warp_n must be 128") + self.warp_n = warp_n + self.mem_n = warp_n + + @staticmethod + def _ceil_to(value: int, divisor: int) -> int: + return (value + divisor - 1) // divisor * divisor + + def _pack_weight_codes(self, codes: torch.Tensor) -> torch.Tensor: + n, k = codes.shape + weight = codes.to(torch.int32).reshape( + n // self.mem_n, + self.mem_n // (self.n_pack_size * self.num_n_lanes * self.reg_n), + self.n_pack_size, + self.num_n_lanes, + self.reg_n, + k // self.mem_k, + 1, + 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() + shifts = torch.arange(0, 32, 4, dtype=torch.int32, device=codes.device) + packed = ((weight & 0xF) << shifts).sum(dim=-1, dtype=torch.int32) + return packed.view(torch.int8).view(n, k // 2) + + def _pack_scale_codes(self, scales: torch.Tensor) -> torch.Tensor: + n, num_groups = scales.shape + scale = scales.view(n // self.warp_n, 1, 4, 4, 8, num_groups // 2, 2) + return scale.permute(0, 5, 1, 4, 3, 2, 6).contiguous().view(num_groups, n) + + def _unpack_weight_codes(self, qweight: torch.Tensor) -> torch.Tensor: + n, packed_k = qweight.shape + k = packed_k * 2 + packed = ( + qweight.contiguous() + .view(torch.int32) + .reshape( + n // self.mem_n, + k // self.mem_k, + 1, + self.mem_n // (self.n_pack_size * self.num_n_lanes * self.reg_n), + self.num_n_lanes, + self.num_k_lanes, + self.n_pack_size, + self.k_pack_size, + self.reg_n, + ) + ) + shifts = torch.arange(0, 32, 4, dtype=torch.int32, device=qweight.device) + weight = (packed.unsqueeze(-1) >> shifts) & 0xF + return weight.permute(0, 3, 6, 4, 8, 1, 2, 7, 5, 9).contiguous().view(n, k).to(torch.uint8) + + def _unpack_scale_codes(self, wscales: torch.Tensor) -> torch.Tensor: + num_groups, n = wscales.shape + scale = wscales.reshape(n // self.warp_n, num_groups // 2, 1, 8, 4, 4, 2) + return scale.permute(0, 2, 5, 4, 3, 1, 6).contiguous().view(n, num_groups) + + def pack_residual(self, weight: torch.Tensor, group_size: int = 32) -> PackedMXFP4: + """Quantize and physically pack a logical ``[N, K]`` residual.""" + + if not isinstance(weight, torch.Tensor) or weight.ndim != 2 or not weight.is_floating_point(): + raise ValueError("weight must be a 2D floating-point torch.Tensor") + if not bool(torch.isfinite(weight).all()): + raise ValueError("weight must contain only finite values") + if group_size != 32: + raise ValueError("group_size must be 32") + + n, k = weight.shape + if n == 0 or k == 0: + raise ValueError("weight dimensions must be non-empty") + n_padded = self._ceil_to(n, self.mem_n) + k_group_padded = self._ceil_to(k, group_size) + k_padded = self._ceil_to(k_group_padded, self.mem_k * self.num_k_unrolls) + + qdq, shared_exponent, _ = quant_mx_rceil( + weight, + bits=4, + group_size=group_size, + data_type="mx_fp4e2m1", + ) + num_logical_groups = k_group_padded // group_size + scales = torch.exp2(shared_exponent.reshape(n, num_logical_groups).to(torch.float32)) + + grouped_qdq = torch.zeros((n, k_group_padded), dtype=qdq.dtype, device=qdq.device) + grouped_qdq[:, :k] = qdq + logical_codes = encode_e2m1(grouped_qdq.reshape(n, num_logical_groups, group_size), scales) + + padded_codes = torch.zeros((n_padded, k_padded), dtype=torch.uint8, device=weight.device) + padded_codes[:n, :k_group_padded] = logical_codes.reshape(n, k_group_padded) + padded_scales = torch.full((n_padded, k_padded // group_size), 127, dtype=torch.uint8, device=weight.device) + padded_scales[:n, :num_logical_groups] = encode_ue8m0(scales) + + return PackedMXFP4( + qweight=self._pack_weight_codes(padded_codes), + wscales=self._pack_scale_codes(padded_scales), + logical_shape=(n, k), + padded_shape=(n_padded, k_padded), + ) + + def unpack_residual( + self, + qweight: torch.Tensor, + wscales: torch.Tensor, + logical_shape: tuple[int, int], + dtype: torch.dtype = torch.float32, + ) -> torch.Tensor: + """Inverse the physical layout and dequantize the logical residual.""" + + if not isinstance(qweight, torch.Tensor) or qweight.dtype != torch.int8: + raise ValueError("qweight must be a torch.int8 tensor") + if qweight.ndim != 2 or qweight.shape[0] == 0 or qweight.shape[0] % self.mem_n: + raise ValueError("qweight shape must be [Npad, Kpad/2] with Npad divisible by 128") + if qweight.shape[1] == 0 or (qweight.shape[1] * 2) % (self.mem_k * self.num_k_unrolls): + raise ValueError("qweight shape must have Kpad divisible by 128") + n_padded, packed_k = qweight.shape + k_padded = packed_k * 2 + expected_scale_shape = (k_padded // 32, n_padded) + if not isinstance(wscales, torch.Tensor) or wscales.dtype != torch.uint8: + raise ValueError("wscales must be a torch.uint8 tensor") + if tuple(wscales.shape) != expected_scale_shape: + raise ValueError(f"wscales shape must be {expected_scale_shape}") + if wscales.device != qweight.device: + raise ValueError("qweight and wscales must be on the same device") + if ( + not isinstance(logical_shape, tuple) + or len(logical_shape) != 2 + or any(isinstance(value, bool) or not isinstance(value, int) for value in logical_shape) + or any(value <= 0 for value in logical_shape) + or logical_shape[0] > n_padded + or logical_shape[1] > k_padded + ): + raise ValueError("logical_shape must be a positive (N, K) tuple within the padded shape") + if dtype not in _SUPPORTED_DECODE_DTYPES: + raise ValueError("dtype must be one of torch.float16, torch.bfloat16, torch.float32, or torch.float64") + weight_codes = self._unpack_weight_codes(qweight) + scale_codes = self._unpack_scale_codes(wscales) + scales = decode_ue8m0(scale_codes) + dequantized = decode_e2m1(weight_codes.reshape(n_padded, k_padded // 32, 32), scales, dtype=dtype).reshape( + n_padded, k_padded + ) + n, k = logical_shape + return dequantized[:n, :k].contiguous() + + +def _broadcast_scales(scales: torch.Tensor, shape: torch.Size, *, device: torch.device) -> torch.Tensor: + """Apply the codec scale layout contract. + + A scalar scale is global. A scale tensor with one fewer dimension than the + values is group-aligned by appending a trailing singleton. A scale tensor + with equal rank uses ordinary PyTorch broadcasting. Other ranks are invalid. + """ + + if not isinstance(scales, torch.Tensor): + raise ValueError("scales must be a torch.Tensor") + if not scales.is_floating_point(): + raise ValueError("scales must have a floating-point dtype") + if scales.device != device: + raise ValueError(f"scales must be on device {device}, got {scales.device}") + if not bool(torch.isfinite(scales).all()) or not bool((scales > 0).all()): + raise ValueError("scales must contain only positive finite values") + tensor_ndim = len(shape) + if scales.ndim == 0: + aligned_scales = scales + layout = "global" + elif scales.ndim == tensor_ndim - 1: + aligned_scales = scales.unsqueeze(-1) + layout = "group-aligned" + elif scales.ndim == tensor_ndim: + aligned_scales = scales + layout = "ordinary-broadcast" + else: + raise ValueError( + f"scales rank {scales.ndim} is invalid for tensor rank {tensor_ndim}; expected a scalar, " + f"rank {tensor_ndim - 1} for group alignment, or rank {tensor_ndim} for ordinary broadcasting" + ) + try: + return torch.broadcast_to(aligned_scales, shape) + except RuntimeError as exc: + raise ValueError( + f"scales shape {tuple(scales.shape)} with {layout} layout cannot broadcast to tensor shape {tuple(shape)}" + ) from exc + + +def _validate_codes(codes: torch.Tensor, *, maximum: int) -> None: + if not isinstance(codes, torch.Tensor): + raise ValueError("codes must be a torch.Tensor") + if codes.dtype == torch.bool or codes.is_floating_point() or codes.is_complex(): + raise ValueError("codes must have an integer dtype") + if codes.numel() and (int(codes.min()) < 0 or int(codes.max()) > maximum): + raise ValueError(f"codes must be in the range 0..{maximum}") + + +def encode_e2m1(values: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: + """Encode raw values as logical E2M1 codes after dividing by ``scales``.""" + + if not isinstance(values, torch.Tensor) or not values.is_floating_point(): + raise ValueError("values must be a floating-point torch.Tensor") + if not bool(torch.isfinite(values).all()): + raise ValueError("values must contain only finite values") + expanded_scales = _broadcast_scales(scales, values.shape, device=values.device) + normalized = values.to(torch.float32) / expanded_scales.to(torch.float32) + invalid = torch.isnan(normalized) + if bool(invalid.any()): + fallback = (values.to(torch.float64) / expanded_scales.to(torch.float64)).clamp(min=-6.0, max=6.0) + normalized = torch.where(invalid, fallback.to(torch.float32), normalized) + normalized = normalized.clamp(min=-6.0, max=6.0) + quantized = quant_element(normalized, ebits=2, mbits=3, max_norm=6.0) + + magnitudes = torch.tensor(_E2M1_MAGNITUDES, dtype=quantized.dtype, device=quantized.device) + magnitude_codes = torch.searchsorted(magnitudes, torch.abs(quantized)).to(torch.uint8) + sign_codes = torch.signbit(normalized).to(torch.uint8) << 3 + return magnitude_codes | sign_codes + + +def decode_e2m1(codes: torch.Tensor, scales: torch.Tensor, dtype: torch.dtype = torch.float32) -> torch.Tensor: + """Decode logical E2M1 codes and multiply by ``scales``.""" + + _validate_codes(codes, maximum=15) + if dtype not in _SUPPORTED_DECODE_DTYPES: + raise ValueError("dtype must be one of torch.float16, torch.bfloat16, torch.float32, or torch.float64") + expanded_scales = _broadcast_scales(scales, codes.shape, device=codes.device) + codebook = torch.tensor( + (*_E2M1_MAGNITUDES, *(-value for value in _E2M1_MAGNITUDES)), + dtype=dtype, + device=codes.device, + ) + return codebook[codes.to(torch.long)] * expanded_scales.to(dtype) + + +def encode_ue8m0(scales: torch.Tensor) -> torch.Tensor: + """Encode scales using an unsigned exponent with bias 127. + + Non-finite and nonpositive scales use code 127, which decodes to 1.0. + """ + + if not isinstance(scales, torch.Tensor) or not scales.is_floating_point(): + raise ValueError("scales must be a floating-point torch.Tensor") + valid = torch.isfinite(scales) & (scales > 0) + safe_scales = torch.where(valid, scales, torch.ones_like(scales)) + exponents = torch.ceil(torch.log2(safe_scales)).clamp(min=-127, max=127) + codes = (exponents + 127).to(torch.uint8) + return torch.where(valid, codes, torch.full_like(codes, 127)) + + +def decode_ue8m0(codes: torch.Tensor) -> torch.Tensor: + """Decode UE8M0 codes to float32 powers of two. + + Code 255 remains decodable for Nunchaku runtime parity, but + :func:`encode_ue8m0` clamps valid exponents to 127 and never emits it. + """ + + _validate_codes(codes, maximum=255) + exponents = codes.to(torch.int16) - 127 + return torch.exp2(exponents.to(torch.float32)) + + +def pack_nibbles(codes: torch.Tensor) -> torch.Tensor: + """Pack logical codes along the last dimension, low nibble first.""" + + _validate_codes(codes, maximum=15) + if codes.ndim == 0: + raise ValueError("codes must have at least one dimension") + codes = codes.to(torch.uint8) + if codes.shape[-1] % 2: + padding = torch.zeros((*codes.shape[:-1], 1), dtype=torch.uint8, device=codes.device) + codes = torch.cat((codes, padding), dim=-1) + pairs = codes.reshape(*codes.shape[:-1], codes.shape[-1] // 2, 2) + return (pairs[..., 0] | (pairs[..., 1] << 4)).contiguous() + + +def unpack_nibbles(packed: torch.Tensor, logical_count: int | None = None) -> torch.Tensor: + """Unpack low-first nibbles along the last dimension.""" + + _validate_codes(packed, maximum=255) + if packed.ndim == 0: + raise ValueError("packed must have at least one dimension") + capacity = packed.shape[-1] * 2 + if logical_count is not None: + if isinstance(logical_count, bool) or not isinstance(logical_count, int): + raise ValueError("logical_count must be an integer or None") + if logical_count < 0 or logical_count > capacity: + raise ValueError(f"logical_count must be between 0 and {capacity}, got {logical_count}") + packed = packed.to(torch.uint8) + unpacked = torch.stack((packed & 0x0F, packed >> 4), dim=-1).reshape(*packed.shape[:-1], capacity) + if logical_count is not None: + unpacked = unpacked[..., :logical_count] + return unpacked.contiguous() diff --git a/auto_round/export/svdquant_w4a16.py b/auto_round/export/svdquant_w4a16.py new file mode 100644 index 0000000000..b82a2dfb4c --- /dev/null +++ b/auto_round/export/svdquant_w4a16.py @@ -0,0 +1,314 @@ +# Copyright (c) 2026 Intel Corporation +# +# 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. + +"""AutoRound-owned Nunchaku W4A16 AdaNorm tensor codec.""" + +from __future__ import annotations + +import sys +from dataclasses import dataclass + +import torch + + +@dataclass(frozen=True) +class PackedW4A16: + """Physical W4A16 tensors and the metadata needed to invert their layout.""" + + qweight: torch.Tensor + wscales: torch.Tensor + wzeros: torch.Tensor + bias: torch.Tensor + dtype: torch.dtype + logical_shape: tuple[int, int] + splits: int + group_size: int + + +def _require_little_endian() -> None: + if sys.byteorder != "little": + raise ValueError("Nunchaku W4A16 packing requires a little-endian host") + + +def _pack_code_rows(codes: torch.Tensor) -> torch.Tensor: + rows, in_features = codes.shape + packed16 = codes.view(-1, 4, 8) + packed16 = packed16[:, 0] | (packed16[:, 1] << 4) | (packed16[:, 2] << 8) | (packed16[:, 3] << 12) + packed16 = ( + packed16.view(rows // 4, 4, in_features // 64, 16) + .permute(0, 2, 1, 3) + .reshape(rows // 4, in_features) + .to(torch.int16) + ) + return packed16.view(torch.int32) + + +def _effective_chunk_rows(chunk_rows: int | None, out_features: int) -> int: + if chunk_rows is None: + return out_features + if isinstance(chunk_rows, bool) or not isinstance(chunk_rows, int) or chunk_rows <= 0 or chunk_rows % 4: + raise ValueError("chunk_rows must be None or a positive multiple of 4") + return min(chunk_rows, out_features) + + +def _validate_inputs( + weight: torch.Tensor, + scale: torch.Tensor, + bias: torch.Tensor | None, + splits: int, + group_size: int, +) -> tuple[int, int, int, torch.Tensor]: + if ( + not isinstance(weight, torch.Tensor) + or weight.ndim != 2 + or weight.dtype + not in ( + torch.bfloat16, + torch.float16, + ) + ): + raise ValueError("weight must be a BF16 or FP16 tensor with shape [O, K]") + if not bool(torch.isfinite(weight).all()): + raise ValueError("weight must contain only finite values") + if isinstance(group_size, bool) or not isinstance(group_size, int) or group_size != 64: + raise ValueError("group_size must be 64") + if isinstance(splits, bool) or not isinstance(splits, int) or splits not in (3, 6): + raise ValueError("splits must be 3 or 6") + out_features, in_features = weight.shape + if out_features == 0 or in_features == 0: + raise ValueError("weight dimensions must be non-empty") + if in_features % group_size: + raise ValueError("K must be divisible by group_size 64") + if out_features % splits: + raise ValueError("O must be divisible by splits") + if out_features % 4: + raise ValueError("O must be divisible by 4") + num_groups = in_features // group_size + if num_groups % 16: + raise ValueError("runtime G=K/64 must be divisible by 16") + if not isinstance(scale, torch.Tensor) or not scale.is_floating_point(): + raise ValueError("scale must be a floating-point tensor") + if scale.dtype != weight.dtype or scale.dtype not in (torch.bfloat16, torch.float16): + raise ValueError("scale dtype must exactly match weight dtype and be BF16 or FP16") + if scale.device != weight.device: + raise ValueError("scale must be on the same device as weight") + if tuple(scale.shape) not in ((out_features, num_groups), (out_features, 1, num_groups, 1)): + raise ValueError(f"scale shape must be [O, G] or [O, 1, G, 1], got {tuple(scale.shape)}") + if not bool(torch.isfinite(scale).all()) or not bool((scale > 0).all()): + raise ValueError("scale must contain only positive finite values") + if bias is not None: + if not isinstance(bias, torch.Tensor) or not bias.is_floating_point(): + raise ValueError("bias must be a floating-point tensor") + if bias.dtype != weight.dtype or bias.dtype not in (torch.bfloat16, torch.float16): + raise ValueError("bias dtype must exactly match weight dtype and be BF16 or FP16") + if tuple(bias.shape) != (out_features,) or bias.device != weight.device or not bool(torch.isfinite(bias).all()): + raise ValueError("bias must be a finite floating-point [O] tensor on the weight device") + return out_features, in_features, num_groups, scale.reshape(out_features, num_groups) + + +def pack_adanorm_w4a16( + weight: torch.Tensor, + scale: torch.Tensor, + bias: torch.Tensor | None = None, + splits: int = 3, + group_size: int = 64, + chunk_rows: int | None = 256, +) -> PackedW4A16: + """Pack pre-scaled signed INT4 AdaNorm weights for Nunchaku W4A16.""" + + _require_little_endian() + out_features, in_features, num_groups, logical_scale = _validate_inputs(weight, scale, bias, splits, group_size) + rows_per_chunk = _effective_chunk_rows(chunk_rows, out_features) + channels_per_field = out_features // splits + qweight = torch.empty((out_features // 4, in_features // 2), dtype=torch.int32, device=weight.device) + channel_scale = torch.empty((out_features, num_groups), dtype=weight.dtype, device=weight.device) + channel_bias = torch.zeros(out_features, dtype=weight.dtype, device=weight.device) + for start in range(0, out_features, rows_per_chunk): + end = min(start + rows_per_chunk, out_features) + output_rows = torch.arange(start, end, device=weight.device) + source_rows = output_rows.remainder(splits) * channels_per_field + torch.div( + output_rows, splits, rounding_mode="floor" + ) + weight_rows = weight.index_select(0, source_rows) + scale_rows = logical_scale.index_select(0, source_rows) + normalized = weight_rows.float().reshape(end - start, num_groups, group_size) + normalized.div_(scale_rows.float().unsqueeze(-1)).round_() + if not bool(((normalized >= -7) & (normalized <= 7)).all()): + raise ValueError("quantized weight must be in [-7, 7]") + codes = normalized.reshape(end - start, in_features).to(torch.int32).add_(7) + qweight[start // 4 : end // 4] = _pack_code_rows(codes) + channel_scale[start:end] = scale_rows + if bias is not None: + channel_bias[start:end] = bias.index_select(0, source_rows) + del codes, normalized, scale_rows, weight_rows + channel_bias = channel_bias.reshape(out_features // splits, splits) + identity_fields = sorted({1, splits - 2}) + identity_before = channel_bias[:, identity_fields].clone() + channel_bias[:, identity_fields] += 1 + if bool((channel_bias[:, identity_fields] == identity_before).any()): + raise ValueError(f"AdaNorm bias identity offset +1 must change the stored {weight.dtype} value") + channel_bias = channel_bias.reshape(out_features) + wscales = channel_scale.t().contiguous() + wzeros = (-7 * channel_scale).t().contiguous() + for name, tensor in (("wscales", wscales), ("wzeros", wzeros), ("bias", channel_bias)): + if not bool(torch.isfinite(tensor).all()): + raise ValueError(f"{name} must remain finite in {weight.dtype} after packing arithmetic") + return PackedW4A16( + qweight=qweight, + wscales=wscales, + wzeros=wzeros, + bias=channel_bias, + dtype=weight.dtype, + logical_shape=(out_features, in_features), + splits=splits, + group_size=group_size, + ) + + +def _validate_packed_w4a16(packed: PackedW4A16) -> tuple[int, int, int]: + _require_little_endian() + if not isinstance(packed, PackedW4A16): + raise ValueError("packed must be a PackedW4A16 payload") + if packed.dtype not in (torch.bfloat16, torch.float16): + raise ValueError("payload dtype must be torch.bfloat16 or torch.float16") + if ( + not isinstance(packed.logical_shape, tuple) + or len(packed.logical_shape) != 2 + or any(isinstance(value, bool) or not isinstance(value, int) or value <= 0 for value in packed.logical_shape) + ): + raise ValueError("logical_shape must be a positive (O, K) integer tuple") + out_features, in_features = packed.logical_shape + if isinstance(packed.splits, bool) or not isinstance(packed.splits, int) or packed.splits not in (3, 6): + raise ValueError("splits must be 3 or 6") + if isinstance(packed.group_size, bool) or not isinstance(packed.group_size, int) or packed.group_size != 64: + raise ValueError("group_size must be 64") + if out_features % packed.splits or out_features % 4 or in_features % packed.group_size: + raise ValueError("logical_shape is incompatible with splits, output packing, or group_size") + num_groups = in_features // packed.group_size + if num_groups % 16: + raise ValueError("logical_shape must produce a runtime group count divisible by 16") + if not isinstance(packed.qweight, torch.Tensor) or packed.qweight.dtype != torch.int32: + raise ValueError("qweight dtype must be torch.int32") + if tuple(packed.qweight.shape) != (out_features // 4, in_features // 2): + raise ValueError("qweight shape is inconsistent with logical_shape") + expected_shapes = { + "wscales": (num_groups, out_features), + "wzeros": (num_groups, out_features), + "bias": (out_features,), + } + for name, expected_shape in expected_shapes.items(): + tensor = getattr(packed, name) + if not isinstance(tensor, torch.Tensor) or tensor.dtype != packed.dtype: + raise ValueError(f"{name} dtype must exactly match payload dtype") + if tuple(tensor.shape) != expected_shape: + raise ValueError(f"{name} shape must be {expected_shape}") + if tensor.device != packed.qweight.device: + raise ValueError("all packed tensors must be on the same device") + for name in ("wscales", "wzeros", "bias"): + if not bool(torch.isfinite(getattr(packed, name)).all()): + raise ValueError(f"{name} must contain only finite values") + if not bool((packed.wscales > 0).all()): + raise ValueError("wscales must contain only positive values") + if not torch.equal(packed.wzeros, packed.wscales * -7): + raise ValueError("wzeros must equal -7 * wscales in payload dtype") + return out_features, in_features, num_groups + + +def unpack_adanorm_w4a16(packed: PackedW4A16) -> torch.Tensor: + """Recover channel-major signed INT4 values from a packed payload.""" + + out_features, in_features, _ = _validate_packed_w4a16(packed) + packed16 = packed.qweight.contiguous().view(torch.int16).reshape(out_features // 4, in_features) + packed16 = packed16.view(out_features // 4, in_features // 64, 4, 16).permute(0, 2, 1, 3).contiguous().view(-1, 8) + shifts = torch.arange(0, 16, 4, dtype=torch.int16, device=packed.qweight.device) + codes = ((packed16.unsqueeze(1) >> shifts.view(1, 4, 1)) & 0xF).reshape(out_features, in_features) + if not bool((codes <= 14).all()): + raise ValueError("qweight contains codes outside [0, 14]") + return (codes - 7).to(torch.int8) + + +def dequantize_adanorm_w4a16(packed: PackedW4A16) -> torch.Tensor: + """Decode the packed weights into their QDQ values in runtime channel order.""" + + signed = unpack_adanorm_w4a16(packed).to(torch.float32) + out_features, in_features = signed.shape + scales = packed.wscales.t().reshape(out_features, in_features // packed.group_size, 1).float() + return (signed.reshape(out_features, -1, packed.group_size) * scales).reshape_as(signed).to(packed.dtype) + + +def _rtn_scale_bounds(dtype: torch.dtype, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]: + zero = torch.tensor(0, dtype=dtype, device=device) + smallest = torch.nextafter(zero, torch.tensor(1, dtype=dtype, device=device)) + largest = torch.tensor(torch.finfo(dtype).max / 7, dtype=dtype, device=device) + while not bool(torch.isfinite(largest * -7)): + largest = torch.nextafter(largest, zero) + return smallest, largest + + +def quantize_adanorm_w4a16_rtn( + weight: torch.Tensor, + bias: torch.Tensor | None = None, + splits: int = 3, + group_size: int = 64, + chunk_rows: int | None = 256, +) -> PackedW4A16: + """Compute symmetric per-row-group RTN scales and pack the result. + + An all-zero group uses scale one. Its signed values and QDQ values remain + exactly zero, while the emitted scale continues to satisfy the runtime's + positive-scale contract. + """ + + _require_little_endian() + if ( + not isinstance(weight, torch.Tensor) + or weight.ndim != 2 + or weight.dtype + not in ( + torch.bfloat16, + torch.float16, + ) + ): + raise ValueError("weight must be a BF16 or FP16 tensor with shape [O, K]") + if isinstance(group_size, bool) or not isinstance(group_size, int) or group_size != 64: + raise ValueError("group_size must be 64") + out_features, in_features = weight.shape + if out_features == 0 or in_features == 0: + raise ValueError("weight dimensions must be non-empty") + if in_features % group_size: + raise ValueError("K must be divisible by group_size 64") + rows_per_chunk = _effective_chunk_rows(chunk_rows, out_features) + num_groups = in_features // group_size + scales = torch.empty((out_features, num_groups), dtype=weight.dtype, device=weight.device) + smallest_scale, largest_scale = _rtn_scale_bounds(weight.dtype, weight.device) + for start in range(0, out_features, rows_per_chunk): + end = min(start + rows_per_chunk, out_features) + grouped = weight[start:end].float().reshape(end - start, num_groups, group_size) + absmax = grouped.abs().amax(dim=-1) + chunk_scales = torch.where(absmax == 0, torch.ones_like(absmax), absmax / 7).to(weight.dtype) + chunk_scales = torch.where((absmax > 0) & (chunk_scales <= 0), smallest_scale, chunk_scales) + scales[start:end] = torch.minimum(chunk_scales, largest_scale) + try: + return pack_adanorm_w4a16( + weight, + scales, + bias=bias, + splits=splits, + group_size=group_size, + chunk_rows=chunk_rows, + ) + except ValueError as exc: + if "quantized weight must be in [-7, 7]" in str(exc): + raise ValueError("RTN scale cannot represent the weight in signed range [-7, 7]") from exc + raise diff --git a/test/test_cpu/export/test_svdquant_mxfp4.py b/test/test_cpu/export/test_svdquant_mxfp4.py new file mode 100644 index 0000000000..299a8c4908 --- /dev/null +++ b/test/test_cpu/export/test_svdquant_mxfp4.py @@ -0,0 +1,55 @@ +# Copyright (c) 2026 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import hashlib + +import pytest +import torch + +from auto_round.data_type.mxfp import quant_mx_rceil +from auto_round.export.svdquant_mxfp4 import ( + NunchakuMXFP4Packer, + pack_lowrank_weight, + unpack_lowrank_weight, +) + + +@pytest.mark.parametrize("down,shape", [(True, (3, 17)), (False, (17, 3))]) +def test_lowrank_pack_roundtrips_and_matches_fixed_layout(down, shape): + logical = torch.arange(shape[0] * shape[1], dtype=torch.float16).reshape(shape) + expected_hash = { + True: "f62c895e44a7139fc942941b1244857d65143dfe52ad852e8847339aa6119029", + False: "0690f5c24d1ca25ad4f7714d9ce6b626b792ac89e46986d7a73dd0f327b163b4", + }[down] + + packed = pack_lowrank_weight(logical, down=down) + unpacked = unpack_lowrank_weight(packed, down=down) + + assert packed.shape == (128, 16) + torch.testing.assert_close(unpacked[: shape[0], : shape[1]], logical) + assert hashlib.sha256(bytes(packed.view(torch.uint8).flatten().tolist())).hexdigest() == expected_hash + + +@pytest.mark.parametrize("shape", [(128, 128), (7, 65)]) +def test_nunchaku_mxfp4_pack_roundtrip_matches_autoround_rceil_qdq(shape): + weight = torch.randn(shape, generator=torch.Generator().manual_seed(20260713)) * 3 + expected, _, _ = quant_mx_rceil(weight, bits=4, group_size=32, data_type="mx_fp4e2m1") + packer = NunchakuMXFP4Packer() + + packed = packer.pack_residual(weight) + actual = packer.unpack_residual(packed.qweight, packed.wscales, packed.logical_shape) + + assert packed.padded_shape[0] % 128 == 0 + assert packed.padded_shape[1] % 128 == 0 + torch.testing.assert_close(actual, expected) diff --git a/test/test_cpu/export/test_svdquant_w4a16.py b/test/test_cpu/export/test_svdquant_w4a16.py new file mode 100644 index 0000000000..b75fdc46ce --- /dev/null +++ b/test/test_cpu/export/test_svdquant_w4a16.py @@ -0,0 +1,56 @@ +# Copyright (c) 2026 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +from auto_round.export.svdquant_w4a16 import ( + dequantize_adanorm_w4a16, + pack_adanorm_w4a16, + quantize_adanorm_w4a16_rtn, + unpack_adanorm_w4a16, +) + + +def _representable_fixture(dtype=torch.float16): + rows = torch.arange(12).reshape(12, 1) + columns = torch.arange(1024).reshape(1, 1024) + signed = ((rows * 5 + columns * 3) % 15 - 7).to(torch.float32) + scales = (torch.arange(12 * 16).reshape(12, 16) % 13 + 1).to(dtype) / 8 + weight = (signed.reshape(12, 16, 64) * scales.unsqueeze(-1)).reshape(12, 1024).to(dtype) + return weight, scales, signed.to(torch.int8) + + +def test_adanorm_w4a16_pack_uses_runtime_layout_and_roundtrips_codes(): + weight, scales, signed = _representable_fixture() + + packed = pack_adanorm_w4a16(weight, scales, splits=3) + + assert packed.qweight.shape == (3, 512) + assert packed.wscales.shape == (16, 12) + assert packed.wzeros.shape == (16, 12) + expected_codes = signed.reshape(3, 4, 1024).permute(1, 0, 2).reshape(12, 1024) + expected_weight = weight.reshape(3, 4, 1024).permute(1, 0, 2).reshape(12, 1024) + assert torch.equal(unpack_adanorm_w4a16(packed), expected_codes) + torch.testing.assert_close(dequantize_adanorm_w4a16(packed), expected_weight) + + +def test_adanorm_w4a16_rtn_emits_finite_runtime_payload(): + weight = torch.randn(12, 1024, generator=torch.Generator().manual_seed(7), dtype=torch.bfloat16) + + packed = quantize_adanorm_w4a16_rtn(weight, splits=3) + + assert packed.qweight.dtype == torch.int32 + assert torch.isfinite(packed.wscales).all() + assert torch.isfinite(packed.wzeros).all() + assert torch.isfinite(dequantize_adanorm_w4a16(packed)).all() From 95837219f50ed2739927da472eb7c7cf1771f237 Mon Sep 17 00:00:00 2001 From: changwangss Date: Tue, 28 Jul 2026 16:55:00 +0800 Subject: [PATCH 6/9] feat: add Nunchaku SVDQuant export Signed-off-by: changwangss --- .../algorithms/transforms/svdquant/apply.py | 1 + auto_round/compressors/base.py | 11 +- auto_round/compressors/diffusion_mixin.py | 37 +- .../export/svdquant_adapters/__init__.py | 75 +++ auto_round/export/svdquant_adapters/flux.py | 625 ++++++++++++++++++ auto_round/export/svdquant_nunchaku.py | 591 +++++++++++++++++ auto_round/formats.py | 151 +++++ auto_round/utils/common.py | 1 + .../export/test_svdquant_flux_adapter.py | 561 ++++++++++++++++ .../export/test_svdquant_nunchaku_export.py | 590 +++++++++++++++++ .../export/test_svdquant_nunchaku_format.py | 241 +++++++ 11 files changed, 2879 insertions(+), 5 deletions(-) create mode 100644 auto_round/export/svdquant_adapters/__init__.py create mode 100644 auto_round/export/svdquant_adapters/flux.py create mode 100644 auto_round/export/svdquant_nunchaku.py create mode 100644 test/test_cpu/export/test_svdquant_flux_adapter.py create mode 100644 test/test_cpu/export/test_svdquant_nunchaku_export.py create mode 100644 test/test_cpu/export/test_svdquant_nunchaku_format.py diff --git a/auto_round/algorithms/transforms/svdquant/apply.py b/auto_round/algorithms/transforms/svdquant/apply.py index f24086665e..539ba61c16 100644 --- a/auto_round/algorithms/transforms/svdquant/apply.py +++ b/auto_round/algorithms/transforms/svdquant/apply.py @@ -158,6 +158,7 @@ def prepare_run(self, composer=None) -> None: self._block_groups.clear() if self.model is None: return + self.model._autoround_svdquant_model_adapter = self.config.model_adapter or "auto" for block_name in self._configured_block_names: block = self.model.get_submodule(block_name) self._block_groups[block_name] = discover_svdquant_groups(block, self._is_target) diff --git a/auto_round/compressors/base.py b/auto_round/compressors/base.py index f0060b54c8..8836391365 100644 --- a/auto_round/compressors/base.py +++ b/auto_round/compressors/base.py @@ -1305,10 +1305,19 @@ def _adjust_immediate_packing_and_saving(self): return formats = getattr(self, "formats", []) + if any(format.requires_full_model_export for format in formats): + self.compress_context.is_immediate_packing = False + self.compress_context.is_immediate_saving = False + has_single_gguf_format = len(formats) == 1 and formats[0].is_gguf() # GGUF supports per-block / per-layer immediate packing even when # full-model in-place rewriting is disabled by outside-block layers. - if len(formats) == 1 and not formats[0].is_fake() and (self.inplace or has_single_gguf_format): + if ( + len(formats) == 1 + and not formats[0].is_fake() + and not formats[0].requires_full_model_export + and (self.inplace or has_single_gguf_format) + ): self.compress_context.is_immediate_packing = True if self.has_qlayer_outside_block and self.need_calib and not has_single_gguf_format: diff --git a/auto_round/compressors/diffusion_mixin.py b/auto_round/compressors/diffusion_mixin.py index aaf25d7845..8273a86843 100644 --- a/auto_round/compressors/diffusion_mixin.py +++ b/auto_round/compressors/diffusion_mixin.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import inspect +import json import os from typing import Any, Optional, Union @@ -29,6 +30,28 @@ from auto_round.utils.model import rename_weights_files +def _rewrite_svdquant_nunchaku_pipeline_index(output_dir: str, component_names: list[str]) -> None: + """Point quantized Diffusers components at their Nunchaku runtime classes.""" + from safetensors import safe_open + + model_index_path = os.path.join(output_dir, "model_index.json") + with open(model_index_path, encoding="utf-8") as file: + model_index = json.load(file) + + for name in component_names: + weight_path = os.path.join(output_dir, name, "diffusion_pytorch_model.safetensors") + with safe_open(weight_path, framework="pt", device="cpu") as file: + metadata = file.metadata() or {} + model_class = metadata.get("model_class") + if not model_class: + raise ValueError(f"{weight_path} is missing required safetensors metadata 'model_class'") + model_index[name] = ["nunchaku", model_class] + + with open(model_index_path, "w", encoding="utf-8") as file: + json.dump(model_index, file, indent=2, sort_keys=True) + file.write("\n") + + class DiffusionMixin: """Diffusion-specific functionality mixin. @@ -486,7 +509,7 @@ def quantize(self) -> tuple[torch.nn.Module, dict]: def save_quantized( self, output_dir: Optional[str] = None, - format: Union[str, list] = "auto_round", + format: Optional[Union[str, list]] = None, inplace: bool = True, return_folders: bool = False, **kwargs, @@ -517,12 +540,15 @@ def save_quantized( has_multiple_quantized_transformers = bool(quantized_transformers) # Handle multi-format (convert string to list if needed) - _format = format + _format = format if format is not None else getattr(self, "formats", None) or "auto_round" if isinstance(_format, str): from auto_round.formats import get_formats _format = get_formats(_format, self) + is_svdquant_nunchaku = any(item.format_name == "svdquant_nunchaku" for item in _format) + quantized_component_names = [] + for name in pipe.components.keys(): val = getattr(pipe, name) sub_module_path = ( @@ -556,6 +582,7 @@ def save_quantized( self.model_context.model._autoround_pipeline_subfolder = saved_subfolder self.model_context.model = saved_model self.layer_config = saved_lc + quantized_component_names.append(name) elif val is self.model_context.model: # Save primary quantized transformer saved_immediate_saving = self.compress_context.is_immediate_saving @@ -574,6 +601,7 @@ def save_quantized( self.compress_context.is_immediate_saving = saved_immediate_saving if saved_subfolder is not None: self.model_context.model._autoround_pipeline_subfolder = saved_subfolder + quantized_component_names.append(name) elif val is not None and hasattr(val, "save_pretrained"): val.save_pretrained(sub_module_path) continue @@ -589,12 +617,13 @@ def save_quantized( pipe.config.save_pretrained(output_dir) else: # FrozenDict / plain dict — write model_index.json manually - import json - model_index_path = os.path.join(output_dir, "model_index.json") with open(model_index_path, "w", encoding="utf-8") as f: f.write(json.dumps(dict(pipe.config), indent=2, sort_keys=True) + "\n") + if is_svdquant_nunchaku: + _rewrite_svdquant_nunchaku_pipeline_index(output_dir, quantized_component_names) + if return_folders: return compressed_model, folders return compressed_model diff --git a/auto_round/export/svdquant_adapters/__init__.py b/auto_round/export/svdquant_adapters/__init__.py new file mode 100644 index 0000000000..e20874ead2 --- /dev/null +++ b/auto_round/export/svdquant_adapters/__init__.py @@ -0,0 +1,75 @@ +# Copyright (c) 2026 Intel Corporation +# +# 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. + +"""Architecture adapters for SVDQuant Nunchaku export.""" + +from __future__ import annotations + +from collections.abc import Mapping + +import torch + +from auto_round.export.svdquant_nunchaku import IdentitySVDQuantModelAdapter + +from auto_round.export.svdquant_adapters.flux import ( + FLUX_TOP_LEVEL_TENSOR_KEYS, + FLUX_SVDQUANT_TARGET_MODULES, + FluxSVDQuantNunchakuAdapter, + flux_onefile_tensor_count, +) + + +def _model_config(model: torch.nn.Module) -> dict: + config = getattr(model, "config", None) + if isinstance(config, Mapping): + return dict(config) + to_dict = getattr(config, "to_dict", None) + if callable(to_dict): + value = to_dict() + if isinstance(value, Mapping): + return dict(value) + return {} + + +def resolve_svdquant_model_adapter( + name: str, + model: torch.nn.Module, + *, + decomposition_device: str | torch.device = "cpu", +): + """Resolve a registered architecture adapter without runtime dependencies.""" + + normalized = name.strip().lower() + if normalized not in {"auto", "identity", "flux"}: + raise ValueError(f"unknown SVDQuant model adapter {name!r}; expected auto, identity, or flux") + config = _model_config(model) + class_name = str(config.get("_class_name", type(model).__name__)).lower() + if normalized == "auto": + normalized = "flux" if "fluxtransformer" in class_name else "identity" + if normalized == "flux": + return FluxSVDQuantNunchakuAdapter( + config=config or None, + decomposition_device=decomposition_device, + require_complete_model=True, + ) + return IdentitySVDQuantModelAdapter() + + +__all__ = [ + "FLUX_TOP_LEVEL_TENSOR_KEYS", + "FLUX_SVDQUANT_TARGET_MODULES", + "FluxSVDQuantNunchakuAdapter", + "flux_onefile_tensor_count", + "resolve_svdquant_model_adapter", +] diff --git a/auto_round/export/svdquant_adapters/flux.py b/auto_round/export/svdquant_adapters/flux.py new file mode 100644 index 0000000000..e95f625f82 --- /dev/null +++ b/auto_round/export/svdquant_adapters/flux.py @@ -0,0 +1,625 @@ +# Copyright (c) 2026 Intel Corporation +# +# 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. + +"""FLUX model mapping for runtime-loadable SVDQuant Nunchaku artifacts.""" + +from __future__ import annotations + +import json +import os +import re +from collections.abc import Iterable, Mapping +from dataclasses import dataclass +from typing import Any + +import torch + +from auto_round.export.svdquant_nunchaku import SourceLinearRecord, SVDQuantExportRecord +from auto_round.export.svdquant_w4a16 import quantize_adanorm_w4a16_rtn + +_BLOCK_RE = re.compile(r"^(transformer_blocks|single_transformer_blocks)\.(\d+)\.(.+)$") +_DOUBLE_DIRECT = { + "attn.to_out.0": "out_proj", + "attn.to_add_out": "out_proj_context", + "ff.net.0.proj": "mlp_fc1", + "ff.net.2": "mlp_fc2", + "ff.net.2.linear": "mlp_fc2", + "ff_context.net.0.proj": "mlp_context_fc1", + "ff_context.net.2": "mlp_context_fc2", + "ff_context.net.2.linear": "mlp_context_fc2", +} +_DOUBLE_FUSED = { + "qkv_proj": ("attn.to_q", "attn.to_k", "attn.to_v"), + "qkv_proj_context": ("attn.add_q_proj", "attn.add_k_proj", "attn.add_v_proj"), +} +_SINGLE_FUSED = {"qkv_proj": ("attn.to_q", "attn.to_k", "attn.to_v")} +_RMS_MAP = { + "attn.norm_q": "norm_q", + "attn.norm_k": "norm_k", + "attn.norm_added_q": "norm_added_q", + "attn.norm_added_k": "norm_added_k", +} +_TOP_LEVEL_PREFIXES = ("x_embedder.", "context_embedder.", "time_text_embed.", "norm_out.linear.", "proj_out.") +FLUX_SVDQUANT_TARGET_MODULES = ( + "attn.to_q", + "attn.to_k", + "attn.to_v", + "attn.add_q_proj", + "attn.add_k_proj", + "attn.add_v_proj", + "attn.to_out.0", + "attn.to_add_out", + "ff.net.0.proj", + "ff.net.2", + "ff_context.net.0.proj", + "ff_context.net.2", + "proj_mlp", + "proj_out", +) +FLUX_TOP_LEVEL_TENSOR_KEYS = frozenset( + { + "x_embedder.weight", + "x_embedder.bias", + "context_embedder.weight", + "context_embedder.bias", + "norm_out.linear.weight", + "norm_out.linear.bias", + "proj_out.weight", + "proj_out.bias", + *( + f"time_text_embed.{embedder}_embedder.linear_{linear}.{parameter}" + for embedder in ("timestep", "text", "guidance") + for linear in (1, 2) + for parameter in ("weight", "bias") + ), + } +) + + +def flux_onefile_tensor_count( + num_layers: int, num_single_layers: int, top_level_tensors: int = len(FLUX_TOP_LEVEL_TENSOR_KEYS) +) -> int: + """Return the key count without constructing model-sized tensors.""" + + return num_layers * (8 * 7 + 2 * 4 + 4) + num_single_layers * (4 * 7 + 4 + 2) + top_level_tensors + + +def _config_dict(value: Any) -> dict[str, Any]: + if isinstance(value, Mapping): + return dict(value) + to_dict = getattr(value, "to_dict", None) + if callable(to_dict): + result = to_dict() + if isinstance(result, Mapping): + return dict(result) + if hasattr(value, "__dict__"): + return {key: item for key, item in vars(value).items() if not key.startswith("_")} + raise ValueError("FLUX config must be a mapping or serialize to a JSON object") + + +def _normalize_config_paths(value: Any) -> Any: + if isinstance(value, os.PathLike): + return os.fspath(value) + if isinstance(value, Mapping): + return {key: _normalize_config_paths(item) for key, item in value.items()} + if isinstance(value, list): + return [_normalize_config_paths(item) for item in value] + if isinstance(value, tuple): + return tuple(_normalize_config_paths(item) for item in value) + return value + + +def _effective_weight(source: SourceLinearRecord, device: torch.device) -> torch.Tensor: + if source.residual_weight.ndim != 2 or source.lora_down.ndim != 2 or source.lora_up.ndim != 2: + raise ValueError(f"{source.name} residual and low-rank weights must be 2D") + out_features, in_features = source.residual_weight.shape + rank = source.lora_down.shape[0] + if ( + source.lora_down.shape != (rank, in_features) + or source.lora_up.shape != (out_features, rank) + or source.smooth.shape != (in_features,) + ): + raise ValueError(f"{source.name} source dimensions are inconsistent") + residual = source.residual_weight.detach().to(device=device, dtype=torch.float32) + up = source.lora_up.detach().to(device=device, dtype=torch.float32) + down = source.lora_down.detach().to(device=device, dtype=torch.float32) + smooth = source.smooth.detach().to(device=device, dtype=torch.float32) + return (residual + up @ down) * smooth.reshape(1, -1) + + +def _has_shared_input_decomposition(sources: tuple[SourceLinearRecord, ...]) -> bool: + first = sources[0] + return all( + source.scheme == first.scheme + and torch.equal(source.lora_down, first.lora_down) + and torch.equal(source.smooth, first.smooth) + and torch.equal(source.smooth_orig, first.smooth_orig) + for source in sources[1:] + ) + + +def _concatenate_shared_sources( + prefix: str, + sources: tuple[SourceLinearRecord, ...], +) -> SVDQuantExportRecord: + first = sources[0] + biases = [source.bias for source in sources] + if any(bias is None for bias in biases) and not all(bias is None for bias in biases): + raise ValueError(f"{prefix} fused sources must either all have bias or all omit bias") + bias = None if biases[0] is None else torch.cat(biases).detach().cpu().contiguous() + return SVDQuantExportRecord( + prefix=prefix, + residual_weight=torch.cat([source.residual_weight for source in sources]).detach().cpu().contiguous(), + lora_down=first.lora_down.detach().cpu().contiguous(), + lora_up=torch.cat([source.lora_up for source in sources]).detach().cpu().contiguous(), + smooth=first.smooth.detach().cpu().contiguous(), + smooth_orig=first.smooth_orig.detach().cpu().contiguous(), + bias=bias, + scheme=first.scheme, + sources=sources, + ) + + +def _decompose( + weight: torch.Tensor, + *, + rank: int, + template: SourceLinearRecord, + prefix: str, + sources: tuple[SourceLinearRecord, ...], + bias: torch.Tensor | None, +) -> SVDQuantExportRecord: + if rank > min(weight.shape): + raise ValueError(f"{prefix} configured rank={rank} exceeds fused dimensions {tuple(weight.shape)}") + if not bool(torch.isfinite(weight).all()): + raise ValueError(f"{prefix} effective weight contains non-finite values") + u, singular_values, vh = torch.linalg.svd(weight, full_matrices=False) + up = u[:, :rank] * singular_values[:rank] + down = vh[:rank] + residual = weight - up @ down + weight_dtype = template.residual_weight.dtype + low_rank_dtype = template.lora_down.dtype + in_features = weight.shape[1] + record = SVDQuantExportRecord( + prefix=prefix, + residual_weight=residual.to(dtype=weight_dtype).cpu().contiguous(), + lora_down=down.to(dtype=low_rank_dtype).cpu().contiguous(), + lora_up=up.to(dtype=template.lora_up.dtype).cpu().contiguous(), + smooth=torch.ones(in_features, dtype=template.smooth.dtype).cpu(), + smooth_orig=torch.ones(in_features, dtype=template.smooth_orig.dtype).cpu(), + bias=None if bias is None else bias.detach().to(dtype=weight_dtype).cpu().contiguous(), + scheme=template.scheme, + sources=sources, + ) + del u, singular_values, vh, up, down, residual, weight + return record + + +@dataclass +class FluxSVDQuantNunchakuAdapter: + """Map Diffusers FLUX modules to the Nunchaku one-file tensor schema.""" + + config: Mapping[str, Any] | None = None + decomposition_device: str | torch.device = "cpu" + require_complete_model: bool = True + + def __post_init__(self) -> None: + try: + self.decomposition_device = torch.device(self.decomposition_device) + except (RuntimeError, TypeError) as exc: + raise ValueError(f"invalid FLUX decomposition_device {self.decomposition_device!r}") from exc + if self.decomposition_device.type not in ("cpu", "cuda"): + raise ValueError("FLUX decomposition_device must be CPU or CUDA") + if self.decomposition_device.type == "cuda": + if not torch.cuda.is_available(): + raise ValueError("FLUX decomposition_device requests CUDA, but CUDA is not available") + device_count = torch.cuda.device_count() + index = self.decomposition_device.index + if index is not None and (index < 0 or index >= device_count): + raise ValueError( + f"FLUX decomposition_device index {index} is invalid for CUDA device_count={device_count}" + ) + if self.config is not None: + self.config = _config_dict(self.config) + if not isinstance(self.require_complete_model, bool): + raise ValueError("require_complete_model must be a bool") + + def _resolved_config(self, model: torch.nn.Module) -> dict[str, Any]: + value = self.config if self.config is not None else getattr(model, "config", None) + if value is None: + raise ValueError("FLUX export requires explicit config or model.config") + config = _normalize_config_paths(_config_dict(value)) + try: + json.dumps(config) + except (TypeError, ValueError) as exc: + raise ValueError("FLUX config must be JSON serializable") from exc + return config + + def metadata(self, model: torch.nn.Module, rank: int) -> Mapping[str, str]: + return { + "model_class": "NunchakuFluxTransformer2dModel", + "config": json.dumps(self._resolved_config(model), sort_keys=True), + "format": "pt", + "comfy_config": "{}", + } + + @staticmethod + def _direct(source: SourceLinearRecord, prefix: str) -> SVDQuantExportRecord: + return SVDQuantExportRecord( + prefix=prefix, + residual_weight=source.residual_weight, + lora_down=source.lora_down, + lora_up=source.lora_up, + smooth=source.smooth, + smooth_orig=source.smooth_orig, + bias=source.bias, + scheme=source.scheme, + sources=(source,), + ) + + def _fuse(self, prefix: str, sources: tuple[SourceLinearRecord, ...], rank: int) -> SVDQuantExportRecord: + def operation() -> SVDQuantExportRecord: + if _has_shared_input_decomposition(sources): + return _concatenate_shared_sources(prefix, sources) + effective = [_effective_weight(source, self.decomposition_device) for source in sources] + input_dims = {weight.shape[1] for weight in effective} + if len(input_dims) != 1: + raise ValueError(f"{prefix} fused sources have incompatible input dimensions {sorted(input_dims)}") + weight = torch.cat(effective, dim=0) + biases = [source.bias for source in sources] + if any(bias is None for bias in biases) and not all(bias is None for bias in biases): + raise ValueError(f"{prefix} fused sources must either all have bias or all omit bias") + bias = None if biases[0] is None else torch.cat([item.to(self.decomposition_device) for item in biases]) + return _decompose(weight, rank=rank, template=sources[0], prefix=prefix, sources=sources, bias=bias) + + try: + return operation() + finally: + self._clear_decomposition_cache() + + def _split_single_proj_out( + self, model: torch.nn.Module, block_prefix: str, source: SourceLinearRecord, rank: int + ) -> tuple[SVDQuantExportRecord, SVDQuantExportRecord]: + def operation() -> tuple[SVDQuantExportRecord, SVDQuantExportRecord]: + config = self._resolved_config(model) + heads, head_dim = config.get("num_attention_heads"), config.get("attention_head_dim") + in_features = source.residual_weight.shape[1] + inner_dim = ( + heads * head_dim + if isinstance(heads, int) and isinstance(head_dim, int) + else source.residual_weight.shape[0] + ) + if inner_dim <= 0 or inner_dim >= in_features: + raise ValueError( + f"{block_prefix}.proj_out cannot split input columns at inner_dim={inner_dim} " + f"for shape {tuple(source.residual_weight.shape)}" + ) + out_proj = SVDQuantExportRecord( + prefix=f"{block_prefix}.out_proj", + residual_weight=source.residual_weight[:, :inner_dim].contiguous(), + lora_down=source.lora_down[:, :inner_dim].contiguous(), + lora_up=source.lora_up, + smooth=source.smooth[:inner_dim].contiguous(), + smooth_orig=source.smooth_orig[:inner_dim].contiguous(), + bias=None, + scheme=source.scheme, + sources=(source,), + ) + mlp_fc2 = SVDQuantExportRecord( + prefix=f"{block_prefix}.mlp_fc2", + residual_weight=source.residual_weight[:, inner_dim:].contiguous(), + lora_down=source.lora_down[:, inner_dim:].contiguous(), + lora_up=source.lora_up, + smooth=source.smooth[inner_dim:].contiguous(), + smooth_orig=source.smooth_orig[inner_dim:].contiguous(), + bias=source.bias, + scheme=source.scheme, + sources=(source,), + ) + return out_proj, mlp_fc2 + + try: + return operation() + finally: + self._clear_decomposition_cache() + + def _clear_decomposition_cache(self) -> None: + if self.decomposition_device.type == "cuda": + torch.cuda.empty_cache() + + def map_modules( + self, model: torch.nn.Module, records: Iterable[SourceLinearRecord] + ) -> Iterable[SVDQuantExportRecord]: + records = tuple(records) + by_block: dict[tuple[str, int], dict[str, SourceLinearRecord]] = {} + for record in records: + match = _BLOCK_RE.match(record.name) + if match is None: + raise ValueError(f"unrecognized FLUX SVDQuant source {record.name!r}") + family, index, local_name = match.groups() + local = by_block.setdefault((family, int(index)), {}) + if local_name in local: + raise ValueError(f"duplicate FLUX source {record.name!r}") + local[local_name] = record + + if self.require_complete_model: + config = self._resolved_config(model) + try: + expected_indices = { + "transformer_blocks": set(range(int(config["num_layers"]))), + "single_transformer_blocks": set(range(int(config["num_single_layers"]))), + } + except (KeyError, TypeError, ValueError) as exc: + raise ValueError("complete FLUX export requires integer num_layers and num_single_layers") from exc + for family, expected in expected_indices.items(): + actual = {index for candidate, index in by_block if candidate == family} + if actual != expected: + raise ValueError( + f"complete FLUX export {family} indices mismatch: expected {sorted(expected)}, " + f"got {sorted(actual)}" + ) + double_fixed = {name for names in _DOUBLE_FUSED.values() for name in names} | { + "attn.to_out.0", + "attn.to_add_out", + "ff.net.0.proj", + "ff_context.net.0.proj", + } + for (family, index), local in by_block.items(): + names = set(local) + if family == "single_transformer_blocks": + required = {"attn.to_q", "attn.to_k", "attn.to_v", "proj_mlp", "proj_out"} + missing, extra = required - names, names - required + else: + missing = double_fixed - names + extra = ( + names + - set(_DOUBLE_DIRECT) + - {name for fused_names in _DOUBLE_FUSED.values() for name in fused_names} + ) + for stem in ("ff.net.2", "ff_context.net.2"): + variants = {stem, f"{stem}.linear"} + present = names & variants + if len(present) != 1: + missing.add(f"exactly one of {sorted(variants)}") + if missing or extra: + raise ValueError( + f"complete FLUX export rejects {family}.{index} coverage: " + f"missing={sorted(missing)}, extras={sorted(extra)}" + ) + + output: list[SVDQuantExportRecord] = [] + for (family, index), local in sorted(by_block.items()): + block_prefix = f"{family}.{index}" + rank = next(iter(local.values())).lora_down.shape[0] + if family == "transformer_blocks": + for target, names in _DOUBLE_FUSED.items(): + present = tuple(local[name] for name in names if name in local) + if present and len(present) != len(names): + missing = [name for name in names if name not in local] + raise ValueError(f"{block_prefix}.{target} missing fused sources {missing}") + if present: + output.append(self._fuse(f"{block_prefix}.{target}", present, rank)) + for source_name, target in _DOUBLE_DIRECT.items(): + if source_name in local: + output.append(self._direct(local[source_name], f"{block_prefix}.{target}")) + else: + names = _SINGLE_FUSED["qkv_proj"] + present = tuple(local[name] for name in names if name in local) + if present and len(present) != len(names): + missing = [name for name in names if name not in local] + raise ValueError(f"{block_prefix}.qkv_proj missing fused sources {missing}") + if present: + output.append(self._fuse(f"{block_prefix}.qkv_proj", present, rank)) + if "proj_mlp" in local: + output.append(self._direct(local["proj_mlp"], f"{block_prefix}.mlp_fc1")) + if "proj_out" in local: + source = local["proj_out"] + output.extend(self._split_single_proj_out(model, block_prefix, source, rank)) + return output + + def validate_records( + self, sources: tuple[SourceLinearRecord, ...], records: tuple[SVDQuantExportRecord, ...] + ) -> None: + ranks = {source.lora_down.shape[0] for source in sources} + if len(ranks) != 1: + raise ValueError(f"FLUX source ranks must agree, got {sorted(ranks)}") + if self.require_complete_model: + # Config was already resolved during mapping for single blocks; retain model-independent + # structural checks here and perform exact source coverage in map_modules via cached names. + families: dict[str, set[int]] = {"transformer_blocks": set(), "single_transformer_blocks": set()} + for source in sources: + match = _BLOCK_RE.match(source.name) + assert match is not None + families[match.group(1)].add(int(match.group(2))) + for family, indices in families.items(): + if indices and indices != set(range(max(indices) + 1)): + raise ValueError(f"complete FLUX export rejects gaps in {family}: {sorted(indices)}") + prefixes = [record.prefix for record in records] + if len(prefixes) != len(set(prefixes)): + raise ValueError("FLUX adapter produced duplicate logical record prefixes") + + @staticmethod + def _module_map(model: torch.nn.Module) -> dict[str, torch.nn.Module]: + return dict(model.named_modules()) + + def extra_tensors(self, model: torch.nn.Module) -> Mapping[str, torch.Tensor]: + modules = self._module_map(model) + tensors: dict[str, torch.Tensor] = {} + block_indices: dict[str, set[int]] = {"transformer_blocks": set(), "single_transformer_blocks": set()} + for name in modules: + match = _BLOCK_RE.match(name) + if match: + block_indices[match.group(1)].add(int(match.group(2))) + + if self.require_complete_model: + config = self._resolved_config(model) + for family, config_name in ( + ("transformer_blocks", "num_layers"), + ("single_transformer_blocks", "num_single_layers"), + ): + expected = set(range(int(config[config_name]))) + if block_indices[family] != expected: + raise ValueError( + f"complete FLUX extras {family} indices mismatch: expected {sorted(expected)}, " + f"got {sorted(block_indices[family])}" + ) + + for family, indices in block_indices.items(): + for index in sorted(indices): + block = f"{family}.{index}" + adanorms = ( + (("norm1.linear", 6), ("norm1_context.linear", 6)) + if family == "transformer_blocks" + else (("norm.linear", 3),) + ) + for local_name, splits in adanorms: + module = modules.get(f"{block}.{local_name}") + if module is None: + if self.require_complete_model: + raise ValueError(f"complete FLUX export missing {block}.{local_name}") + continue + weight = getattr(module, "weight", None) + bias = getattr(module, "bias", None) + if not isinstance(weight, torch.Tensor): + raise ValueError(f"{block}.{local_name}.weight must be a tensor") + packed = quantize_adanorm_w4a16_rtn( + weight.detach().to(device="cpu", dtype=torch.bfloat16), + bias=None if bias is None else bias.detach().to(device="cpu", dtype=torch.bfloat16), + splits=splits, + group_size=64, + ) + prefix = f"{block}.{local_name}" + tensors.update( + { + f"{prefix}.qweight": packed.qweight, + f"{prefix}.wscales": packed.wscales, + f"{prefix}.wzeros": packed.wzeros, + f"{prefix}.bias": packed.bias, + } + ) + for source_name, target_name in _RMS_MAP.items(): + module = modules.get(f"{block}.{source_name}") + weight = None if module is None else getattr(module, "weight", None) + if weight is None: + required_rms = family == "transformer_blocks" or source_name in ("attn.norm_q", "attn.norm_k") + if self.require_complete_model and required_rms: + raise ValueError(f"complete FLUX export missing {block}.{source_name}.weight") + continue + tensors[f"{block}.{target_name}.weight"] = weight.detach().to(torch.bfloat16).cpu() + + passthrough = { + name: parameter for name, parameter in model.named_parameters() if name.startswith(_TOP_LEVEL_PREFIXES) + } + if self.require_complete_model: + missing_top_level = FLUX_TOP_LEVEL_TENSOR_KEYS - passthrough.keys() + if missing_top_level: + raise ValueError(f"complete FLUX model is missing top-level parameters: {sorted(missing_top_level)}") + extra_top_level = passthrough.keys() - FLUX_TOP_LEVEL_TENSOR_KEYS + if extra_top_level: + raise ValueError(f"complete FLUX model has unexpected top-level parameters: {sorted(extra_top_level)}") + for name in FLUX_TOP_LEVEL_TENSOR_KEYS & passthrough.keys(): + tensors[name] = passthrough[name].detach().to(torch.bfloat16).cpu() + return {key: value.contiguous() for key, value in tensors.items()} + + def validate(self, tensors: Mapping[str, torch.Tensor], metadata: Mapping[str, str]) -> None: + if metadata.get("model_class") != "NunchakuFluxTransformer2dModel": + raise ValueError("FLUX metadata has incorrect model_class") + if metadata.get("format") != "pt" or metadata.get("comfy_config") != "{}": + raise ValueError("FLUX metadata requires format='pt' and empty comfy_config") + try: + config = json.loads(metadata["config"]) + except (KeyError, json.JSONDecodeError) as exc: + raise ValueError("FLUX metadata config must be a JSON object") from exc + if not isinstance(config, dict): + raise ValueError("FLUX metadata config must be a JSON object") + required: set[str] = set() + if self.require_complete_model: + try: + num_layers = int(config["num_layers"]) + num_single_layers = int(config["num_single_layers"]) + except (KeyError, TypeError, ValueError) as exc: + raise ValueError("complete FLUX metadata requires layer counts") from exc + top_level_keys = { + key for key in tensors if not key.startswith(("transformer_blocks.", "single_transformer_blocks.")) + } + missing_top_level = FLUX_TOP_LEVEL_TENSOR_KEYS - top_level_keys + if missing_top_level: + raise ValueError(f"complete FLUX artifact is missing top-level tensors: {sorted(missing_top_level)}") + extra_top_level = top_level_keys - FLUX_TOP_LEVEL_TENSOR_KEYS + if extra_top_level: + raise ValueError(f"complete FLUX artifact has unexpected top-level tensors: {sorted(extra_top_level)}") + linear_suffixes = ("qweight", "wscales", "smooth", "smooth_orig", "lora_down", "lora_up", "bias") + for index in range(num_layers): + block = f"transformer_blocks.{index}" + for linear in ( + "qkv_proj", + "qkv_proj_context", + "out_proj", + "out_proj_context", + "mlp_fc1", + "mlp_fc2", + "mlp_context_fc1", + "mlp_context_fc2", + ): + required.update(f"{block}.{linear}.{suffix}" for suffix in linear_suffixes) + for norm in ("norm1.linear", "norm1_context.linear"): + required.update(f"{block}.{norm}.{suffix}" for suffix in ("qweight", "wscales", "wzeros", "bias")) + required.update( + f"{block}.{norm}.weight" for norm in ("norm_q", "norm_k", "norm_added_q", "norm_added_k") + ) + for index in range(num_single_layers): + block = f"single_transformer_blocks.{index}" + for linear in ("qkv_proj", "out_proj", "mlp_fc1", "mlp_fc2"): + required.update(f"{block}.{linear}.{suffix}" for suffix in linear_suffixes) + required.update(f"{block}.norm.linear.{suffix}" for suffix in ("qweight", "wscales", "wzeros", "bias")) + required.update(f"{block}.{norm}.weight" for norm in ("norm_q", "norm_k")) + missing = required - tensors.keys() + if missing: + raise ValueError(f"complete FLUX artifact is missing expected tensors: {sorted(missing)[:5]}") + if num_layers == 19 and num_single_layers == 38 and len(tensors) != flux_onefile_tensor_count(19, 38): + raise ValueError(f"standard FLUX one-file artifact must contain 2604 tensors, got {len(tensors)}") + for key, tensor in tensors.items(): + if tensor.device.type != "cpu" or not tensor.is_contiguous(): + raise ValueError(f"FLUX tensor {key!r} must be contiguous on CPU") + if tensor.is_floating_point() and not bool(torch.isfinite(tensor).all()): + raise ValueError(f"FLUX tensor {key!r} must be finite") + if not key.startswith(("transformer_blocks.", "single_transformer_blocks.")): + if not key.startswith(_TOP_LEVEL_PREFIXES) or tensor.dtype != torch.bfloat16: + raise ValueError(f"FLUX passthrough tensor {key!r} must be in an allowed top-level BF16 family") + is_adanorm = ".norm" in key and any( + marker in key for marker in (".norm.linear.", ".norm1.linear.", ".norm1_context.linear.") + ) + if key.endswith(".qweight"): + expected_dtype = torch.int32 if is_adanorm else torch.int8 + if tensor.dtype != expected_dtype or tensor.ndim != 2: + raise ValueError(f"FLUX qweight {key!r} must be 2D {expected_dtype}") + elif key.endswith(".wscales"): + expected_dtype = torch.bfloat16 if is_adanorm else torch.uint8 + if tensor.dtype != expected_dtype or tensor.ndim != 2: + raise ValueError(f"FLUX wscales {key!r} must be 2D {expected_dtype}") + elif key.endswith(".wzeros"): + if not is_adanorm or tensor.dtype != torch.bfloat16 or tensor.ndim != 2: + raise ValueError(f"FLUX wzeros {key!r} must be a 2D AdaNorm BF16 tensor") + elif key.endswith((".lora_down", ".lora_up")): + if tensor.dtype != torch.bfloat16 or tensor.ndim != 2: + raise ValueError(f"FLUX low-rank tensor {key!r} must be 2D BF16") + elif key.endswith((".smooth", ".smooth_orig", ".bias")): + if tensor.dtype != torch.bfloat16 or tensor.ndim != 1: + raise ValueError(f"FLUX vector tensor {key!r} must be 1D BF16") + elif key.endswith(".weight") and key.startswith(("transformer_blocks.", "single_transformer_blocks.")): + if tensor.dtype != torch.bfloat16 or tensor.ndim != 1: + raise ValueError(f"FLUX RMSNorm tensor {key!r} must be 1D BF16") + + +__all__ = ["FLUX_TOP_LEVEL_TENSOR_KEYS", "FluxSVDQuantNunchakuAdapter", "flux_onefile_tensor_count"] diff --git a/auto_round/export/svdquant_nunchaku.py b/auto_round/export/svdquant_nunchaku.py new file mode 100644 index 0000000000..e493e9ad87 --- /dev/null +++ b/auto_round/export/svdquant_nunchaku.py @@ -0,0 +1,591 @@ +# Copyright (c) 2026 Intel Corporation +# +# 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. + +from __future__ import annotations + +import json +import os +from dataclasses import dataclass, field +from typing import Iterable, Mapping, Protocol + +import torch + +from auto_round.algorithms.transforms.svdquant.wrapper import SVDQuantLinear +from auto_round.export.svdquant_mxfp4 import NunchakuMXFP4Packer, pack_lowrank_weight +from auto_round.wrapper import WrapperWALayer + +_DEPLOYABLE_E2M1_ALIASES = frozenset({"mx_fp", "mx_fp4", "mx_fp4e2m1"}) +NUNCHAKU_WEIGHT_FILENAME = "diffusion_pytorch_model.safetensors" + + +class ResidualTensorProvider(Protocol): + """Provides packed MXFP4 residual tensors for one logical export record.""" + + def tensors_for(self, record: SVDQuantExportRecord) -> Mapping[str, torch.Tensor]: + """Return tensors keyed by suffix, for example ``qweight`` and ``wscales``.""" + + +@dataclass(frozen=True) +class SVDQuantLinearScheme: + """AutoRound-selected weight and activation scheme values.""" + + data_type: str | None + bits: int | None + group_size: int | tuple[int, int] | None + sym: bool | None + act_data_type: str | None + act_bits: int | None + act_group_size: int | tuple[int, int] | None + act_sym: bool | None + act_dynamic: bool | None + + +@dataclass(frozen=True) +class SourceLinearRecord: + """Logical tensors from one source ``SVDQuantLinear`` before adapter mapping.""" + + name: str + residual_weight: torch.Tensor + lora_down: torch.Tensor + lora_up: torch.Tensor + smooth: torch.Tensor + smooth_orig: torch.Tensor + bias: torch.Tensor | None + scheme: SVDQuantLinearScheme + + +@dataclass(frozen=True) +class SVDQuantExportRecord: + """Adapter-selected logical tensors and their exported key mapping. + + Fusion adapters must reconstruct effective source weights, fuse or split those + weights, and recompute the low-rank decomposition at the configured output + rank. Nunchaku metadata stores one rank, so exact rank-sum LoRA fusion is not + representable by this export schema. + """ + + prefix: str + residual_weight: torch.Tensor + lora_down: torch.Tensor + lora_up: torch.Tensor + smooth: torch.Tensor + smooth_orig: torch.Tensor + bias: torch.Tensor | None + scheme: SVDQuantLinearScheme + sources: tuple[SourceLinearRecord, ...] + key_mapping: Mapping[str, str] = field(default_factory=dict) + + +class SVDQuantModelAdapter(Protocol): + """Boundary for model-level mapping and runtime metadata. + + ``map_modules`` receives every logical source together so architecture adapters + can recompose effective weights before sibling fusion or splitting. Its output + must retain source provenance and use the configured source rank. + """ + + def map_modules( + self, model: torch.nn.Module, records: Iterable[SourceLinearRecord] + ) -> Iterable[SVDQuantExportRecord]: ... + + def metadata(self, model: torch.nn.Module, rank: int) -> Mapping[str, str]: ... + + def extra_tensors(self, model: torch.nn.Module) -> Mapping[str, torch.Tensor]: ... + + def validate_records( + self, sources: tuple[SourceLinearRecord, ...], records: tuple[SVDQuantExportRecord, ...] + ) -> None: ... + + def validate(self, tensors: Mapping[str, torch.Tensor], metadata: Mapping[str, str]) -> None: ... + + +class IdentitySVDQuantModelAdapter: + """Map modules unchanged and explicitly identify a generic intermediate.""" + + def map_modules( + self, model: torch.nn.Module, records: Iterable[SourceLinearRecord] + ) -> Iterable[SVDQuantExportRecord]: + return ( + SVDQuantExportRecord( + prefix=record.name or "model", + residual_weight=record.residual_weight, + lora_down=record.lora_down, + lora_up=record.lora_up, + smooth=record.smooth, + smooth_orig=record.smooth_orig, + bias=record.bias, + scheme=record.scheme, + sources=(record,), + ) + for record in records + ) + + def metadata(self, model: torch.nn.Module, rank: int) -> Mapping[str, str]: + return {"artifact_type": "generic_intermediate"} + + def extra_tensors(self, model: torch.nn.Module) -> Mapping[str, torch.Tensor]: + return {} + + def validate_records( + self, sources: tuple[SourceLinearRecord, ...], records: tuple[SVDQuantExportRecord, ...] + ) -> None: + return + + def validate(self, tensors: Mapping[str, torch.Tensor], metadata: Mapping[str, str]) -> None: + if metadata.get("artifact_type") != "generic_intermediate": + raise ValueError("identity adapter output must be marked as a generic intermediate") + + +@dataclass +class SVDQuantExportConfig: + """Strict configuration for SVDQuant Nunchaku serialization.""" + + weight_dtype: str = "fp4_e2m1_all" + activation_dtype: str = "fp4_e2m1_all" + scale_dtype: str = "ue8m0" + group_size: int = 32 + low_rank_dtype: torch.dtype = torch.bfloat16 + debug_unpacked: bool = False + runtime_loadable: bool = False + + def __post_init__(self) -> None: + if self.weight_dtype != "fp4_e2m1_all": + raise ValueError("weight_dtype must be 'fp4_e2m1_all'") + if self.activation_dtype != "fp4_e2m1_all": + raise ValueError("activation_dtype must be 'fp4_e2m1_all'") + if self.scale_dtype != "ue8m0": + raise ValueError("scale_dtype must be 'ue8m0'") + if isinstance(self.group_size, bool) or not isinstance(self.group_size, int) or self.group_size != 32: + raise ValueError("group_size must be 32") + if self.low_rank_dtype not in (torch.bfloat16, torch.float16): + raise ValueError("low_rank_dtype must be torch.bfloat16 or torch.float16") + if not isinstance(self.debug_unpacked, bool): + raise ValueError("debug_unpacked must be a bool") + if not isinstance(self.runtime_loadable, bool): + raise ValueError("runtime_loadable must be a bool") + + def to_quantization_config(self) -> dict: + config = { + "method": "svdquant", + "weight": { + "dtype": self.weight_dtype, + "scale_dtype": self.scale_dtype, + "group_size": self.group_size, + }, + "activation": { + "dtype": self.activation_dtype, + "scale_dtype": self.scale_dtype, + "group_size": self.group_size, + }, + } + return config + + +def pack_nunchaku_16bit_vector(vector: torch.Tensor) -> torch.Tensor: + """Pack and pad one BF16/FP16 vector like Nunchaku's scalar scale path.""" + + if not isinstance(vector, torch.Tensor) or vector.ndim != 1: + raise ValueError("vector must be a 1D torch.Tensor") + if vector.dtype not in (torch.bfloat16, torch.float16): + raise ValueError("vector dtype must be torch.bfloat16 or torch.float16") + if vector.numel() == 0: + raise ValueError("vector must be non-empty") + if not bool(torch.isfinite(vector).all()): + raise ValueError("vector must contain only finite values") + padded_size = NunchakuMXFP4Packer._ceil_to(vector.numel(), 128) + padded = torch.ones(padded_size, dtype=vector.dtype, device=vector.device) + padded[: vector.numel()] = vector + packed = padded.reshape(padded_size // 128, 1, 8, 2, 4, 2, 1) + return packed.permute(0, 6, 1, 2, 4, 3, 5).contiguous().view(-1) + + +def unpack_nunchaku_16bit_vector(vector: torch.Tensor) -> torch.Tensor: + """Invert :func:`pack_nunchaku_16bit_vector`, retaining padding.""" + + if not isinstance(vector, torch.Tensor) or vector.ndim != 1: + raise ValueError("vector must be a 1D torch.Tensor") + if vector.dtype not in (torch.bfloat16, torch.float16): + raise ValueError("vector dtype must be torch.bfloat16 or torch.float16") + if vector.numel() == 0 or vector.numel() % 128: + raise ValueError("packed vector length must be a non-zero multiple of 128") + unpacked = vector.reshape(vector.numel() // 128, 1, 1, 8, 4, 2, 2) + return unpacked.permute(0, 2, 3, 5, 4, 6, 1).contiguous().view(-1) + + +class MXFP4ResidualTensorProvider: + """Quantize and physically pack residual weights as E2M1 with UE8M0 scales.""" + + def __init__(self, group_size: int = 32) -> None: + if isinstance(group_size, bool) or not isinstance(group_size, int) or group_size != 32: + raise ValueError("group_size must be 32") + self.group_size = group_size + self.packer = NunchakuMXFP4Packer() + + def tensors_for(self, record: SVDQuantExportRecord) -> Mapping[str, torch.Tensor]: + weight = record.residual_weight + if not bool(torch.isfinite(weight).all()): + raise ValueError(f"{record.prefix or ''} residual weight must contain only finite values") + packed = self.packer.pack_residual(weight, group_size=self.group_size) + return {"qweight": packed.qweight, "wscales": packed.wscales} + + +def _source_records(model: torch.nn.Module) -> tuple[SourceLinearRecord, ...]: + records = [] + for name, module in model.named_modules(): + if not isinstance(module, SVDQuantLinear): + continue + residual_linear = module.residual_linear + while isinstance(residual_linear, WrapperWALayer): + residual_linear = residual_linear.orig_layer + records.append( + SourceLinearRecord( + name=name, + residual_weight=residual_linear.weight.detach(), + lora_down=module.lora_down.weight.detach(), + lora_up=module.lora_up.weight.detach(), + smooth=module.smooth.detach(), + smooth_orig=getattr(module, "smooth_orig", module.smooth).detach(), + bias=None if residual_linear.bias is None else residual_linear.bias.detach(), + scheme=SVDQuantLinearScheme( + data_type=getattr(residual_linear, "data_type", None), + bits=getattr(residual_linear, "bits", None), + group_size=getattr(residual_linear, "group_size", None), + sym=getattr(residual_linear, "sym", None), + act_data_type=getattr(residual_linear, "act_data_type", None), + act_bits=getattr(residual_linear, "act_bits", None), + act_group_size=getattr(residual_linear, "act_group_size", None), + act_sym=getattr(residual_linear, "act_sym", None), + act_dynamic=getattr(residual_linear, "act_dynamic", None), + ), + ) + ) + if not records: + raise ValueError("No SVDQuantLinear modules found to export.") + return tuple(records) + + +def _validate_selected_scheme(scheme: SVDQuantLinearScheme, prefix: str) -> tuple: + required_weight = { + "data_type": scheme.data_type, + "bits": scheme.bits, + "group_size": scheme.group_size, + "sym": scheme.sym, + } + missing_weight = [name for name, value in required_weight.items() if value is None] + if missing_weight: + raise ValueError(f"{prefix} selected residual scheme is missing required {missing_weight[0]}") + weight_data_type = scheme.data_type.removesuffix("_rceil") if isinstance(scheme.data_type, str) else None + if weight_data_type not in _DEPLOYABLE_E2M1_ALIASES: + raise ValueError( + f"{prefix} residual data_type must be one of {sorted(_DEPLOYABLE_E2M1_ALIASES)}, got {scheme.data_type!r}" + ) + if isinstance(scheme.bits, bool) or not isinstance(scheme.bits, int) or scheme.bits != 4: + raise ValueError(f"{prefix} residual scheme requires bits=4, got {scheme.bits!r}") + if isinstance(scheme.group_size, bool) or not isinstance(scheme.group_size, int) or scheme.group_size != 32: + raise ValueError(f"{prefix} residual scheme requires scalar group_size=32, got {scheme.group_size!r}") + if scheme.sym is not True: + raise ValueError(f"{prefix} residual scheme requires sym=True, got {scheme.sym!r}") + + required_activation = { + "act_data_type": scheme.act_data_type, + "act_bits": scheme.act_bits, + "act_group_size": scheme.act_group_size, + "act_sym": scheme.act_sym, + "act_dynamic": scheme.act_dynamic, + } + missing_activation = [name for name, value in required_activation.items() if value is None] + if missing_activation: + raise ValueError(f"{prefix} selected scheme is missing required activation value {missing_activation[0]}") + activation_data_type = ( + scheme.act_data_type.removesuffix("_rceil") if isinstance(scheme.act_data_type, str) else None + ) + if activation_data_type not in _DEPLOYABLE_E2M1_ALIASES: + raise ValueError( + f"{prefix} activation data_type must be one of {sorted(_DEPLOYABLE_E2M1_ALIASES)}, " + f"got {scheme.act_data_type!r}" + ) + if isinstance(scheme.act_bits, bool) or not isinstance(scheme.act_bits, int) or scheme.act_bits != 4: + raise ValueError(f"{prefix} activation scheme requires activation bits=4, got {scheme.act_bits!r}") + if ( + isinstance(scheme.act_group_size, bool) + or not isinstance(scheme.act_group_size, int) + or scheme.act_group_size != 32 + ): + raise ValueError( + f"{prefix} activation scheme requires activation scalar group_size=32, got {scheme.act_group_size!r}" + ) + if scheme.act_sym is not True: + raise ValueError(f"{prefix} activation scheme requires activation sym=True, got {scheme.act_sym!r}") + if scheme.act_dynamic is not True: + raise ValueError(f"{prefix} activation scheme requires act_dynamic=True, got {scheme.act_dynamic!r}") + return ("mx_fp4e2m1", 4, 32, True, "mx_fp4e2m1", 4, 32, True, True) + + +def _validate_export_record(record: SVDQuantExportRecord, config: SVDQuantExportConfig) -> int: + tensors = (record.residual_weight, record.lora_down, record.lora_up, record.smooth, record.smooth_orig) + if not record.prefix or not isinstance(record.prefix, str): + raise ValueError("export record prefix must be a non-empty string") + if any(not isinstance(tensor, torch.Tensor) or not tensor.is_floating_point() for tensor in tensors): + raise ValueError(f"{record.prefix} logical tensors must be floating-point torch.Tensor values") + if record.residual_weight.ndim != 2 or record.lora_down.ndim != 2 or record.lora_up.ndim != 2: + raise ValueError(f"{record.prefix} residual and low-rank tensors must be 2D") + out_features, in_features = record.residual_weight.shape + rank = record.lora_down.shape[0] + if ( + record.lora_down.shape[1] != in_features + or record.lora_up.shape != (out_features, rank) + or record.smooth.shape != (in_features,) + or record.smooth_orig.shape != (in_features,) + or (record.bias is not None and record.bias.shape != (out_features,)) + ): + raise ValueError(f"{record.prefix} logical tensor shapes are inconsistent") + finite_tensors = tensors + (() if record.bias is None else (record.bias,)) + if any(not bool(torch.isfinite(tensor).all()) for tensor in finite_tensors): + raise ValueError(f"{record.prefix} logical tensors must contain only finite values") + if rank <= 0: + raise ValueError(f"{record.prefix} rank must be positive") + selected_scheme = _validate_selected_scheme(record.scheme, record.prefix) + if not record.sources or any(not isinstance(source, SourceLinearRecord) for source in record.sources): + raise ValueError(f"{record.prefix} export record must retain at least one logical source record") + source_schemes = {_validate_selected_scheme(source.scheme, source.name) for source in record.sources} + if len(source_schemes) != 1 or selected_scheme not in source_schemes: + raise ValueError(f"{record.prefix} adapter sources have incompatible selected quantization schemes") + if config.group_size != selected_scheme[2]: + raise ValueError( + f"{record.prefix} export group_size={config.group_size} disagrees with " + f"selected group_size={selected_scheme[2]}" + ) + return rank + + +def _export_key(record: SVDQuantExportRecord, suffix: str) -> str: + mapped_suffix = record.key_mapping.get(suffix, suffix) + if not isinstance(mapped_suffix, str) or not mapped_suffix: + raise ValueError(f"{record.prefix} key mapping for {suffix!r} must be a non-empty string") + return f"{record.prefix}.{mapped_suffix}" + + +def _validate_packed_residual(payload: Mapping[str, torch.Tensor], record: SVDQuantExportRecord) -> None: + prefix = record.prefix + if set(payload) != {"qweight", "wscales"}: + raise ValueError(f"{prefix} packed residual must contain exactly qweight and wscales") + qweight, wscales = payload["qweight"], payload["wscales"] + if not isinstance(qweight, torch.Tensor) or qweight.dtype != torch.int8 or qweight.ndim != 2: + raise ValueError(f"{prefix} qweight must be a 2D torch.int8 tensor") + out_features, in_features = record.residual_weight.shape + padded_out = NunchakuMXFP4Packer._ceil_to(out_features, 128) + padded_in = NunchakuMXFP4Packer._ceil_to(in_features, 128) + expected_weight_shape = (padded_out, padded_in // 2) + if tuple(qweight.shape) != expected_weight_shape: + raise ValueError(f"{prefix} qweight shape must be {expected_weight_shape}") + expected_scale_shape = (padded_in // 32, padded_out) + if not isinstance(wscales, torch.Tensor) or wscales.dtype != torch.uint8: + raise ValueError(f"{prefix} wscales must be a torch.uint8 tensor") + if tuple(wscales.shape) != expected_scale_shape: + raise ValueError(f"{prefix} wscales shape must be {expected_scale_shape}") + + +def _validate_adapter_provenance( + sources: tuple[SourceLinearRecord, ...], records: tuple[SVDQuantExportRecord, ...] +) -> None: + """Require complete source identity coverage and configured-rank outputs.""" + + source_ids = {id(source) for source in sources} + referenced_ids: set[int] = set() + for record in records: + output_rank = record.lora_down.shape[0] + for source in record.sources: + if id(source) not in source_ids: + raise ValueError(f"{record.prefix} adapter record references a foreign logical source") + if source.lora_down.shape[0] != output_rank: + raise ValueError( + f"{record.prefix} source rank={source.lora_down.shape[0]} must equal output rank={output_rank}; " + "the Nunchaku schema stores one configured rank, so adapters must recompose effective source " + "weights and decompose them at that configured rank; exact rank-sum fusion is unsupported" + ) + referenced_ids.add(id(source)) + missing = [source.name or "" for source in sources if id(source) not in referenced_ids] + if missing: + raise ValueError(f"model adapter dropped logical sources: {missing}") + + +def _prepare_export_records( + model: torch.nn.Module, + config: SVDQuantExportConfig, + adapter: SVDQuantModelAdapter, +) -> tuple[tuple[SourceLinearRecord, ...], tuple[SVDQuantExportRecord, ...], int]: + source_records = _source_records(model) + for source in source_records: + _validate_selected_scheme(source.scheme, source.name or "") + records = tuple(adapter.map_modules(model, source_records)) + if not records: + raise ValueError("model adapter produced no SVDQuant export records") + ranks = {_validate_export_record(record, config) for record in records} + if len(ranks) != 1: + raise ValueError(f"mixed SVDQuant ranks are not supported: {sorted(ranks)}") + _validate_adapter_provenance(source_records, records) + validate_records = getattr(adapter, "validate_records", None) + if validate_records is not None: + validate_records(source_records, records) + return source_records, records, ranks.pop() + + +def _serialize_export_records( + records: tuple[SVDQuantExportRecord, ...], + config: SVDQuantExportConfig, + residual_provider: ResidualTensorProvider, +) -> dict[str, torch.Tensor]: + tensors: dict[str, torch.Tensor] = {} + for record in records: + if not bool((record.smooth > 0).all()) or not bool((record.smooth_orig > 0).all()): + raise ValueError(f"{record.prefix} smooth tensors must contain only positive values") + + # SVDQuantLinear multiplies both branches' input by ``smooth``. Nunchaku + # divides only the residual input by its smooth factor and applies the + # low-rank branch to the original input, so convert between coordinates. + runtime_smooth = record.smooth.to(torch.float64).reciprocal() + runtime_smooth_orig = record.smooth_orig.to(torch.float64).reciprocal() + runtime_lora_down = record.lora_down.to(torch.float64) * record.smooth.to(torch.float64).unsqueeze(0) + high_precision = { + "smooth": pack_nunchaku_16bit_vector(runtime_smooth.to(config.low_rank_dtype)), + "smooth_orig": pack_nunchaku_16bit_vector(runtime_smooth_orig.to(config.low_rank_dtype)), + "lora_down": pack_lowrank_weight(runtime_lora_down.to(config.low_rank_dtype), down=True), + "lora_up": pack_lowrank_weight(record.lora_up.to(config.low_rank_dtype), down=False), + "bias": pack_nunchaku_16bit_vector( + torch.zeros( + record.residual_weight.shape[0], + dtype=config.low_rank_dtype, + device=record.residual_weight.device, + ) + if record.bias is None + else record.bias.to(config.low_rank_dtype) + ), + } + residual_payload = residual_provider.tensors_for(record) + _validate_packed_residual(residual_payload, record) + serialized = {**high_precision, **residual_payload} + if config.debug_unpacked: + serialized["residual.weight"] = record.residual_weight + for suffix, tensor in serialized.items(): + key = _export_key(record, suffix) + if key in tensors: + raise ValueError(f"model adapter produced duplicate tensor key {key!r}") + tensors[key] = tensor.detach().cpu().contiguous() + return tensors + + +def _merge_adapter_tensors( + tensors: dict[str, torch.Tensor], adapter: SVDQuantModelAdapter, model: torch.nn.Module +) -> dict[str, torch.Tensor]: + extra_tensors = getattr(adapter, "extra_tensors", None) + extras = {} if extra_tensors is None else extra_tensors(model) + if not isinstance(extras, Mapping): + raise ValueError("model adapter extra_tensors must return a mapping") + for key, tensor in extras.items(): + if not isinstance(key, str) or not key: + raise ValueError("model adapter extra tensor keys must be non-empty strings") + if key in tensors: + raise ValueError(f"model adapter produced duplicate tensor key {key!r}") + if not isinstance(tensor, torch.Tensor): + raise ValueError(f"model adapter extra tensor {key!r} must be a torch.Tensor") + if tensor.layout != torch.strided or tensor.is_quantized: + raise ValueError(f"model adapter extra tensor {key!r} must be a dense strided tensor") + if tensor.is_complex() or tensor.dtype == torch.bool: + raise ValueError(f"model adapter extra tensor {key!r} has unsupported dtype {tensor.dtype}") + if (tensor.is_floating_point() or tensor.is_complex()) and not bool(torch.isfinite(tensor).all()): + raise ValueError(f"model adapter extra tensor {key!r} must contain only finite values") + tensors[key] = tensor.detach().cpu().contiguous() + return tensors + + +def collect_svdquant_tensors( + model: torch.nn.Module, + *, + config: SVDQuantExportConfig | None = None, + residual_provider: ResidualTensorProvider | None = None, + adapter: SVDQuantModelAdapter | None = None, +) -> dict[str, torch.Tensor]: + """Collect generic SVDQuant tensors from all ``SVDQuantLinear`` modules.""" + + config = config or SVDQuantExportConfig() + residual_provider = residual_provider or MXFP4ResidualTensorProvider(config.group_size) + adapter = adapter or IdentitySVDQuantModelAdapter() + _, records, rank = _prepare_export_records(model, config, adapter) + if config.runtime_loadable: + build_svdquant_metadata(model, config=config, adapter=adapter, rank=rank) + tensors = _serialize_export_records(records, config, residual_provider) + return _merge_adapter_tensors(tensors, adapter, model) + + +def build_svdquant_metadata( + model: torch.nn.Module, + *, + config: SVDQuantExportConfig | None = None, + adapter: SVDQuantModelAdapter | None = None, + rank: int | None = None, +) -> dict[str, str]: + """Build validated string metadata for one SVDQuant artifact.""" + + config = config or SVDQuantExportConfig() + adapter = adapter or IdentitySVDQuantModelAdapter() + if rank is None: + _, _, rank = _prepare_export_records(model, config, adapter) + quantization_config = config.to_quantization_config() + quantization_config["rank"] = rank + metadata = dict(adapter.metadata(model, rank)) + metadata["quantization_config"] = json.dumps(quantization_config, sort_keys=True) + if not metadata or any(not isinstance(key, str) or not isinstance(value, str) for key, value in metadata.items()): + raise ValueError("safetensors metadata keys and values must be strings") + if config.runtime_loadable: + if config.debug_unpacked: + raise ValueError("debug_unpacked output is not runtime-loadable") + if not metadata.get("model_class") or not metadata.get("config"): + raise ValueError("runtime-loadable export requires adapter metadata 'model_class' and serialized 'config'") + try: + serialized_config = json.loads(metadata["config"]) + except (TypeError, json.JSONDecodeError) as exc: + raise ValueError("adapter metadata 'config' must be serialized JSON") from exc + if not isinstance(serialized_config, dict): + raise ValueError("adapter metadata 'config' must serialize a JSON object") + return metadata + + +def save_svdquant_nunchaku_safetensors( + model: torch.nn.Module, + output_path: str, + *, + config: SVDQuantExportConfig | None = None, + residual_provider: ResidualTensorProvider | None = None, + adapter: SVDQuantModelAdapter | None = None, +) -> str: + """Save decomposed SVDQuant tensors to one safetensors file. + + This function is model-family agnostic. It exports every ``SVDQuantLinear`` + it finds and does not import inference runtimes. + """ + + from safetensors.torch import save_file + + config = config or SVDQuantExportConfig() + adapter = adapter or IdentitySVDQuantModelAdapter() + residual_provider = residual_provider or MXFP4ResidualTensorProvider(config.group_size) + _, records, rank = _prepare_export_records(model, config, adapter) + metadata = build_svdquant_metadata(model, config=config, adapter=adapter, rank=rank) + tensors = _serialize_export_records(records, config, residual_provider) + tensors = _merge_adapter_tensors(tensors, adapter, model) + adapter.validate(tensors, metadata) + + os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True) + save_file(tensors, output_path, metadata=metadata) + return output_path diff --git a/auto_round/formats.py b/auto_round/formats.py index 314ceb9cf3..02a42cd74a 100644 --- a/auto_round/formats.py +++ b/auto_round/formats.py @@ -210,6 +210,7 @@ class OutputFormat(ABC): support_schemes: list = [] _format_list: dict[str, OutputFormat] = {} format_name = "base" + requires_full_model_export = False def __init__(self, format: str, ar: BaseCompressor): """Initialize the OutputFormat class.""" @@ -332,6 +333,156 @@ def is_llm_compressor(self) -> bool: return "llm_compressor" in self.output_format or (self.backend is not None and self.backend.is_llm_compressor()) +@OutputFormat.register("svdquant_nunchaku") +class SVDQuantNunchakuFormat(OutputFormat): + support_schemes = ["MXFP4"] + format_name = "svdquant_nunchaku" + requires_full_model_export = True + _e2m1_aliases = frozenset({"mx_fp", "mx_fp4", "mx_fp4e2m1"}) + + def __init__(self, format: str, ar: BaseCompressor): + self.output_format = format + self.backend = None + if not self.is_support_scheme(ar.scheme): + raise ValueError(f"{self.format_name} supports only MXFP4 E2M1 group32 export; got scheme {ar.scheme!r}.") + preset = PRESET_SCHEMES["MXFP4"] + resolved = QuantizationScheme.from_dict( + {name: getattr(ar, name, getattr(preset, name)) for name in QuantizationScheme.get_attributes()} + ) + self.check_scheme_args(resolved) + self._resolved_scheme = resolved + + @classmethod + def check_scheme_args(cls, scheme: QuantizationScheme) -> bool: + rules = { + "data_type": cls._e2m1_aliases, + "bits": 4, + "group_size": 32, + "sym": True, + "act_data_type": cls._e2m1_aliases, + "act_bits": 4, + "act_group_size": 32, + "act_sym": True, + "act_dynamic": True, + } + for name, expected in rules.items(): + actual = getattr(scheme, name, None) + if isinstance(expected, frozenset): + valid = isinstance(actual, str) and actual in expected + expected_text = f"one of {sorted(expected)}" + elif isinstance(expected, int) and not isinstance(expected, bool): + valid = isinstance(actual, int) and not isinstance(actual, bool) and actual == expected + expected_text = repr(expected) + else: + valid = actual is expected + expected_text = repr(expected) + if not valid: + raise ValueError( + f"{cls.format_name} got {name}={actual!r}; expected {name}={expected_text} " + "for Nunchaku E2M1 group32 export." + ) + return True + + def _validate_svd_layer_overrides(self, model: torch.nn.Module, layer_config: dict | None) -> None: + if model is None or not layer_config: + return + from auto_round.algorithms.transforms.svdquant.wrapper import SVDQuantLinear + + for name, module in model.named_modules(): + if not isinstance(module, SVDQuantLinear): + continue + for layer_name in (name, f"{name}.residual_linear"): + if layer_name not in layer_config: + continue + override = layer_config[layer_name] + if isinstance(override, str): + scheme = PRESET_SCHEMES[override.upper()].copy() + elif isinstance(override, QuantizationScheme): + scheme = override.copy() + elif isinstance(override, dict): + values = dict(override) + preset_name = values.pop("scheme", None) + scheme = ( + PRESET_SCHEMES[preset_name.upper()].copy() + if preset_name is not None + else self._resolved_scheme.copy() + ) + scheme.update_from_dict(values) + else: + raise TypeError( + f"Unsupported layer_config value for SVDQuant residual {layer_name!r}: {type(override)}" + ) + try: + self.check_scheme_args(scheme) + except ValueError as exc: + raise ValueError( + f"{self.format_name} layer {layer_name!r} has an incompatible residual scheme: {exc}" + ) from exc + + def check_and_reset_format(self, ar: BaseCompressor) -> None: + self._validate_svd_layer_overrides(getattr(ar, "model", None), getattr(ar, "layer_config", None)) + return None + + def pack_layer(self, *args, **kwargs): + return None + + def save_quantized( + self, + output_dir: str, + model: torch.nn.Module = None, + tokenizer: Callable = None, + layer_config: dict = None, + inplace: bool = True, + device: Union[str, torch.device] = "cpu", + serialization_dict: dict = None, + *, + config=None, + residual_provider=None, + adapter=None, + model_adapter=None, + **kwargs, + ) -> torch.nn.Module: + if output_dir is None: + return model + from auto_round.export.svdquant_nunchaku import ( + NUNCHAKU_WEIGHT_FILENAME, + SVDQuantExportConfig, + save_svdquant_nunchaku_safetensors, + ) + + self._validate_svd_layer_overrides(model, layer_config) + model_adapter = model_adapter or getattr(model, "_autoround_svdquant_model_adapter", "auto") + if isinstance(model_adapter, str): + from auto_round.export.svdquant_adapters import resolve_svdquant_model_adapter + + model_adapter = resolve_svdquant_model_adapter(model_adapter, model, decomposition_device=device) + if model_adapter is not None: + if adapter is not None: + raise TypeError("Pass only one of model_adapter and adapter.") + adapter = model_adapter + from auto_round.export.svdquant_nunchaku import IdentitySVDQuantModelAdapter + + if isinstance(adapter, IdentitySVDQuantModelAdapter): + raise ValueError( + "svdquant_nunchaku requires a runtime model adapter; " + "use a supported architecture or select one with model_adapter" + ) + if config is None: + config = SVDQuantExportConfig(runtime_loadable=True) + elif not config.runtime_loadable: + raise ValueError("svdquant_nunchaku requires config.runtime_loadable=True") + + output_path = os.path.join(os.fspath(output_dir), NUNCHAKU_WEIGHT_FILENAME) + save_svdquant_nunchaku_safetensors( + model, + output_path, + config=config, + residual_provider=residual_provider, + adapter=adapter, + ) + return model + + @OutputFormat.register("fake") class FakeFormat(OutputFormat): support_schemes = None diff --git a/auto_round/utils/common.py b/auto_round/utils/common.py index 4af0291a2c..c23aad8eb2 100644 --- a/auto_round/utils/common.py +++ b/auto_round/utils/common.py @@ -619,6 +619,7 @@ def __init__(self): "auto_round:fp8", "mlx", "auto_round:mlx", + "svdquant_nunchaku", ) self._gguf_format = tuple(sorted(GGUF_CONFIG.keys())) self._support_list = self._support_format + self._gguf_format diff --git a/test/test_cpu/export/test_svdquant_flux_adapter.py b/test/test_cpu/export/test_svdquant_flux_adapter.py new file mode 100644 index 0000000000..eb3e548ad2 --- /dev/null +++ b/test/test_cpu/export/test_svdquant_flux_adapter.py @@ -0,0 +1,561 @@ +import inspect +import json +from pathlib import Path + +import pytest +import torch + +from auto_round.export.svdquant_adapters.flux import ( + FluxSVDQuantNunchakuAdapter, + flux_onefile_tensor_count, +) +from auto_round.export.svdquant_nunchaku import ( + IdentitySVDQuantModelAdapter, + SourceLinearRecord, + SVDQuantExportConfig, + SVDQuantLinearScheme, + collect_svdquant_tensors, + save_svdquant_nunchaku_safetensors, +) + +SCHEME = SVDQuantLinearScheme("mx_fp4", 4, 32, True, "mx_fp4", 4, 32, True, True) + + +def _source(name, out_features=8, in_features=8, rank=2, seed=0, bias=True): + generator = torch.Generator().manual_seed(seed) + return SourceLinearRecord( + name=name, + residual_weight=torch.randn(out_features, in_features, generator=generator), + lora_down=torch.randn(rank, in_features, generator=generator), + lora_up=torch.randn(out_features, rank, generator=generator), + smooth=torch.linspace(0.5, 1.5, in_features), + smooth_orig=torch.linspace(0.75, 1.75, in_features), + bias=torch.randn(out_features, generator=generator) if bias else None, + scheme=SCHEME, + ) + + +def _effective(source): + return (source.residual_weight + source.lora_up @ source.lora_down) * source.smooth + + +def _model(config=None): + model = torch.nn.Module() + model.config = config or { + "num_layers": 1, + "num_single_layers": 1, + "num_attention_heads": 2, + "attention_head_dim": 4, + } + return model + + +def test_double_qkv_reconstructs_effective_sources_at_fixed_rank_in_order(): + sources = tuple( + _source(f"transformer_blocks.0.attn.to_{name}", seed=index + 1) for index, name in enumerate(("q", "k", "v")) + ) + adapter = FluxSVDQuantNunchakuAdapter(require_complete_model=False) + + (record,) = tuple(adapter.map_modules(_model(), sources)) + + assert record.prefix == "transformer_blocks.0.qkv_proj" + assert record.sources == sources + assert record.lora_down.shape[0] == 2 + assert torch.equal(record.smooth, torch.ones(8)) + expected = torch.cat([_effective(source) for source in sources]) + actual = record.residual_weight + record.lora_up @ record.lora_down + torch.testing.assert_close(actual, expected, atol=2e-5, rtol=2e-5) + torch.testing.assert_close(record.bias, torch.cat([source.bias for source in sources])) + assert all( + tensor.device.type == "cpu" and tensor.is_contiguous() + for tensor in (record.residual_weight, record.lora_down, record.lora_up) + ) + + +def test_double_context_qkv_reconstructs_effective_sources_with_runtime_name_and_bias_order(): + sources = tuple( + _source(f"transformer_blocks.0.attn.add_{name}_proj", seed=index + 11) + for index, name in enumerate(("q", "k", "v")) + ) + + (record,) = tuple(FluxSVDQuantNunchakuAdapter(require_complete_model=False).map_modules(_model(), sources)) + + assert record.prefix == "transformer_blocks.0.qkv_proj_context" + expected = torch.cat([_effective(source) for source in sources]) + torch.testing.assert_close( + record.residual_weight + record.lora_up @ record.lora_down, expected, atol=2e-5, rtol=2e-5 + ) + torch.testing.assert_close(record.bias, torch.cat([source.bias for source in sources])) + + +def test_double_qkv_preserves_shared_low_rank_and_smooth_without_redecomposition(): + shared_down = torch.randn(2, 8) + shared_smooth = torch.linspace(0.5, 1.5, 8) + shared_smooth_orig = torch.linspace(0.75, 1.75, 8) + sources = tuple( + _source(f"transformer_blocks.0.attn.to_{name}", seed=index + 21) for index, name in enumerate(("q", "k", "v")) + ) + sources = tuple( + SourceLinearRecord( + name=source.name, + residual_weight=source.residual_weight, + lora_down=shared_down.clone(), + lora_up=source.lora_up, + smooth=shared_smooth.clone(), + smooth_orig=shared_smooth_orig.clone(), + bias=source.bias, + scheme=source.scheme, + ) + for source in sources + ) + + (record,) = tuple(FluxSVDQuantNunchakuAdapter(require_complete_model=False).map_modules(_model(), sources)) + + torch.testing.assert_close(record.residual_weight, torch.cat([source.residual_weight for source in sources])) + torch.testing.assert_close(record.lora_down, shared_down) + torch.testing.assert_close(record.lora_up, torch.cat([source.lora_up for source in sources])) + torch.testing.assert_close(record.smooth, shared_smooth) + torch.testing.assert_close(record.smooth_orig, shared_smooth_orig) + + +def test_single_proj_out_splits_input_columns_and_keeps_bias_only_on_mlp(): + source = _source("single_transformer_blocks.0.proj_out", out_features=8, in_features=16, seed=10) + model = _model({"num_layers": 0, "num_single_layers": 1, "num_attention_heads": 2, "attention_head_dim": 4}) + adapter = FluxSVDQuantNunchakuAdapter(require_complete_model=False) + + out_proj, mlp_fc2 = tuple(adapter.map_modules(model, (source,))) + + assert (out_proj.prefix, mlp_fc2.prefix) == ( + "single_transformer_blocks.0.out_proj", + "single_transformer_blocks.0.mlp_fc2", + ) + expected = _effective(source) + out_effective = (out_proj.residual_weight + out_proj.lora_up @ out_proj.lora_down) * out_proj.smooth + mlp_effective = (mlp_fc2.residual_weight + mlp_fc2.lora_up @ mlp_fc2.lora_down) * mlp_fc2.smooth + torch.testing.assert_close(out_effective, expected[:, :8], atol=2e-5, rtol=2e-5) + torch.testing.assert_close(mlp_effective, expected[:, 8:], atol=2e-5, rtol=2e-5) + torch.testing.assert_close(out_proj.lora_down, source.lora_down[:, :8]) + torch.testing.assert_close(mlp_fc2.lora_down, source.lora_down[:, 8:]) + torch.testing.assert_close(out_proj.smooth, source.smooth[:8]) + torch.testing.assert_close(mlp_fc2.smooth, source.smooth[8:]) + assert out_proj.bias is None + torch.testing.assert_close(mlp_fc2.bias, source.bias) + + +@pytest.mark.parametrize( + ("source_name", "target_name"), + [ + ("transformer_blocks.0.attn.to_out.0", "transformer_blocks.0.out_proj"), + ("transformer_blocks.0.ff.net.2.linear", "transformer_blocks.0.mlp_fc2"), + ("transformer_blocks.0.ff_context.net.0.proj", "transformer_blocks.0.mlp_context_fc1"), + ("single_transformer_blocks.0.proj_mlp", "single_transformer_blocks.0.mlp_fc1"), + ], +) +def test_direct_maps_preserve_logical_records(source_name, target_name): + source = _source(source_name) + (record,) = tuple(FluxSVDQuantNunchakuAdapter(require_complete_model=False).map_modules(_model(), (source,))) + assert record.prefix == target_name + assert record.residual_weight is source.residual_weight + assert record.lora_down is source.lora_down + assert record.smooth is source.smooth + + +def _install(root, path, module): + current = root + parts = path.split(".") + for part in parts[:-1]: + if not hasattr(current, part): + current.add_module(part, torch.nn.Module()) + current = getattr(current, part) + current.add_module(parts[-1], module) + + +def _install_parameter(root, name, tensor): + module_name, parameter_name = name.rsplit(".", 1) + current = root + for part in module_name.split("."): + if not hasattr(current, part): + current.add_module(part, torch.nn.Module()) + current = getattr(current, part) + current.register_parameter(parameter_name, torch.nn.Parameter(tensor)) + + +def test_extra_tensors_pack_adanorm_copy_norms_and_top_level_bf16(): + model = _model() + for name, splits in ( + ("transformer_blocks.0.norm1.linear", 6), + ("transformer_blocks.0.norm1_context.linear", 6), + ("single_transformer_blocks.0.norm.linear", 3), + ): + linear = torch.nn.Linear(1024, 12, bias=True, dtype=torch.bfloat16) + _install(model, name, linear) + for local_name in ("norm_q", "norm_k", "norm_added_q", "norm_added_k"): + norm = torch.nn.Module() + norm.weight = torch.nn.Parameter(torch.randn(8)) + _install(model, f"transformer_blocks.0.attn.{local_name}", norm) + _install(model, "x_embedder", torch.nn.Linear(8, 8)) + _install(model, "unrelated", torch.nn.Linear(8, 8)) + + tensors = FluxSVDQuantNunchakuAdapter(require_complete_model=False).extra_tensors(model) + + for prefix in ( + "transformer_blocks.0.norm1.linear", + "transformer_blocks.0.norm1_context.linear", + "single_transformer_blocks.0.norm.linear", + ): + assert {f"{prefix}.{suffix}" for suffix in ("qweight", "wscales", "wzeros", "bias")} <= tensors.keys() + assert tensors["transformer_blocks.0.norm_added_k.weight"].dtype == torch.bfloat16 + assert tensors["x_embedder.weight"].dtype == torch.bfloat16 + assert not any(key.startswith("unrelated.") for key in tensors) + assert all(tensor.device.type == "cpu" and tensor.is_contiguous() for tensor in tensors.values()) + + +def test_metadata_explicit_config_and_complete_mode_rejects_gaps(): + adapter = FluxSVDQuantNunchakuAdapter(config={"num_layers": 2, "num_single_layers": 0}, require_complete_model=True) + metadata = adapter.metadata(_model(), 2) + assert metadata["model_class"] == "NunchakuFluxTransformer2dModel" + assert json.loads(metadata["config"])["num_layers"] == 2 + assert metadata["format"] == "pt" and metadata["comfy_config"] == "{}" + with pytest.raises(ValueError, match="indices mismatch"): + tuple(adapter.map_modules(_model(), (_source("transformer_blocks.1.attn.to_q"),))) + + +def test_metadata_normalizes_pathlike_config_values(): + adapter = FluxSVDQuantNunchakuAdapter( + config={ + "num_layers": 0, + "num_single_layers": 0, + "_name_or_path": Path("/models/flux/transformer"), + "nested": {"cache": Path("/cache/flux")}, + } + ) + + config = json.loads(adapter.metadata(_model(), 2)["config"]) + + assert config["_name_or_path"] == "/models/flux/transformer" + assert config["nested"]["cache"] == "/cache/flux" + + +def test_decomposition_device_accepts_available_cuda_index(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "device_count", lambda: 2) + + adapter = FluxSVDQuantNunchakuAdapter(decomposition_device="cuda:1", require_complete_model=False) + + assert adapter.decomposition_device == torch.device("cuda:1") + + +@pytest.mark.parametrize( + ("device", "available", "count", "message"), + [ + ("cuda", False, 0, "CUDA.*not available"), + ("cuda:2", True, 2, "index 2.*device_count=2"), + ("mps", True, 2, "must be CPU or CUDA"), + ], +) +def test_decomposition_device_rejects_unavailable_or_unsupported_accelerator( + monkeypatch, device, available, count, message +): + monkeypatch.setattr(torch.cuda, "is_available", lambda: available) + monkeypatch.setattr(torch.cuda, "device_count", lambda: count) + + with pytest.raises(ValueError, match=message): + FluxSVDQuantNunchakuAdapter(decomposition_device=device, require_complete_model=False) + + +def test_cuda_decomposition_groups_return_cpu_and_release_cache_without_requiring_cuda(monkeypatch): + import auto_round.export.svdquant_adapters.flux as flux_module + + requested_devices = [] + empty_cache_calls = [] + original_effective_weight = flux_module._effective_weight + + def effective_weight_on_cpu(source, device): + requested_devices.append(device) + return original_effective_weight(source, torch.device("cpu")) + + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "device_count", lambda: 1) + monkeypatch.setattr(torch.cuda, "empty_cache", lambda: empty_cache_calls.append(True)) + monkeypatch.setattr(flux_module, "_effective_weight", effective_weight_on_cpu) + sources = tuple( + _source(f"transformer_blocks.0.attn.to_{name}", seed=index + 31, bias=False) + for index, name in enumerate(("q", "k", "v")) + ) + + (record,) = tuple( + FluxSVDQuantNunchakuAdapter(decomposition_device="cuda:0", require_complete_model=False).map_modules( + _model(), sources + ) + ) + + assert requested_devices == [torch.device("cuda:0")] * 3 + assert empty_cache_calls == [True] + assert all( + tensor.device.type == "cpu" + for tensor in (record.residual_weight, record.lora_down, record.lora_up, record.smooth) + ) + + +def test_cuda_single_split_avoids_redecomposition_and_releases_cache_once(monkeypatch): + import auto_round.export.svdquant_adapters.flux as flux_module + + requested_devices = [] + empty_cache_calls = [] + original_effective_weight = flux_module._effective_weight + + def effective_weight_on_cpu(source, device): + requested_devices.append(device) + return original_effective_weight(source, torch.device("cpu")) + + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "device_count", lambda: 1) + monkeypatch.setattr(torch.cuda, "empty_cache", lambda: empty_cache_calls.append(True)) + monkeypatch.setattr(flux_module, "_effective_weight", effective_weight_on_cpu) + source = _source("single_transformer_blocks.0.proj_out", out_features=8, in_features=16, seed=37) + + records = tuple( + FluxSVDQuantNunchakuAdapter(decomposition_device="cuda", require_complete_model=False).map_modules( + _model(), (source,) + ) + ) + + assert requested_devices == [] + assert empty_cache_calls == [True] + assert len(records) == 2 + assert all(record.residual_weight.device.type == "cpu" for record in records) + + +def test_cpu_decomposition_does_not_clear_cuda_cache(monkeypatch): + empty_cache_calls = [] + monkeypatch.setattr(torch.cuda, "empty_cache", lambda: empty_cache_calls.append(True)) + sources = tuple( + _source(f"transformer_blocks.0.attn.to_{name}", seed=index + 41) for index, name in enumerate(("q", "k", "v")) + ) + + tuple(FluxSVDQuantNunchakuAdapter(require_complete_model=False).map_modules(_model(), sources)) + + assert empty_cache_calls == [] + + +def test_standard_flux_onefile_count_formula(): + assert flux_onefile_tensor_count(19, 38) == 2604 + + +def test_complete_flux_top_level_keyset_is_exact_known_good_schema(): + from auto_round.export.svdquant_adapters import flux + + expected = { + "x_embedder.weight", + "x_embedder.bias", + "context_embedder.weight", + "context_embedder.bias", + "norm_out.linear.weight", + "norm_out.linear.bias", + "proj_out.weight", + "proj_out.bias", + } + expected.update( + f"time_text_embed.{embedder}_embedder.linear_{linear}.{parameter}" + for embedder in ("timestep", "text", "guidance") + for linear in (1, 2) + for parameter in ("weight", "bias") + ) + + assert flux.FLUX_TOP_LEVEL_TENSOR_KEYS == frozenset(expected) + assert len(flux.FLUX_TOP_LEVEL_TENSOR_KEYS) == 20 + + +def _small_complete_top_level_tensors(): + from auto_round.export.svdquant_adapters.flux import FLUX_TOP_LEVEL_TENSOR_KEYS + + return { + key: torch.ones(2, dtype=torch.bfloat16) if key.endswith(".bias") else torch.ones(2, 2, dtype=torch.bfloat16) + for key in FLUX_TOP_LEVEL_TENSOR_KEYS + } + + +def _small_linear_tensors(prefix): + return { + f"{prefix}.qweight": torch.ones(1, 1, dtype=torch.int8), + f"{prefix}.wscales": torch.ones(1, 1, dtype=torch.uint8), + f"{prefix}.smooth": torch.ones(1, dtype=torch.bfloat16), + f"{prefix}.smooth_orig": torch.ones(1, dtype=torch.bfloat16), + f"{prefix}.lora_down": torch.ones(1, 1, dtype=torch.bfloat16), + f"{prefix}.lora_up": torch.ones(1, 1, dtype=torch.bfloat16), + f"{prefix}.bias": torch.ones(1, dtype=torch.bfloat16), + } + + +def _small_adanorm_tensors(prefix): + return { + f"{prefix}.qweight": torch.ones(1, 1, dtype=torch.int32), + f"{prefix}.wscales": torch.ones(1, 1, dtype=torch.bfloat16), + f"{prefix}.wzeros": torch.ones(1, 1, dtype=torch.bfloat16), + f"{prefix}.bias": torch.ones(1, dtype=torch.bfloat16), + } + + +def _small_standard_complete_tensors(): + tensors = _small_complete_top_level_tensors() + for index in range(19): + block = f"transformer_blocks.{index}" + for linear in ( + "qkv_proj", + "qkv_proj_context", + "out_proj", + "out_proj_context", + "mlp_fc1", + "mlp_fc2", + "mlp_context_fc1", + "mlp_context_fc2", + ): + tensors.update(_small_linear_tensors(f"{block}.{linear}")) + tensors.update(_small_adanorm_tensors(f"{block}.norm1.linear")) + tensors.update(_small_adanorm_tensors(f"{block}.norm1_context.linear")) + for norm in ("norm_q", "norm_k", "norm_added_q", "norm_added_k"): + tensors[f"{block}.{norm}.weight"] = torch.ones(1, dtype=torch.bfloat16) + for index in range(38): + block = f"single_transformer_blocks.{index}" + for linear in ("qkv_proj", "out_proj", "mlp_fc1", "mlp_fc2"): + tensors.update(_small_linear_tensors(f"{block}.{linear}")) + tensors.update(_small_adanorm_tensors(f"{block}.norm.linear")) + for norm in ("norm_q", "norm_k"): + tensors[f"{block}.{norm}.weight"] = torch.ones(1, dtype=torch.bfloat16) + return tensors + + +def test_complete_schema_validation_rejects_missing_top_level_key(): + model = _model({"num_layers": 0, "num_single_layers": 0}) + adapter = FluxSVDQuantNunchakuAdapter(require_complete_model=True) + metadata = adapter.metadata(model, 2) + tensors = _small_complete_top_level_tensors() + tensors.pop("x_embedder.bias") + + with pytest.raises(ValueError, match="missing top-level.*x_embedder.bias"): + adapter.validate(tensors, metadata) + + +def test_complete_schema_validation_accepts_exact_small_key_mapping(): + model = _model({"num_layers": 0, "num_single_layers": 0}) + adapter = FluxSVDQuantNunchakuAdapter(require_complete_model=True) + + adapter.validate(_small_complete_top_level_tensors(), adapter.metadata(model, 2)) + + +def test_standard_complete_schema_validates_2604_tiny_tensors_without_model_allocation(): + model = _model({"num_layers": 19, "num_single_layers": 38}) + adapter = FluxSVDQuantNunchakuAdapter(require_complete_model=True) + tensors = _small_standard_complete_tensors() + + assert len(tensors) == 2604 + adapter.validate(tensors, adapter.metadata(model, 32)) + + +def test_complete_schema_validation_rejects_junk_passthrough_key(): + model = _model({"num_layers": 0, "num_single_layers": 0}) + adapter = FluxSVDQuantNunchakuAdapter(require_complete_model=True) + tensors = _small_complete_top_level_tensors() + tensors["x_embedder.junk"] = torch.ones(2, dtype=torch.bfloat16) + + with pytest.raises(ValueError, match="unexpected top-level.*x_embedder.junk"): + adapter.validate(tensors, adapter.metadata(model, 2)) + + +def test_complete_extra_collection_rejects_junk_passthrough_parameter(): + model = _model({"num_layers": 0, "num_single_layers": 0}) + for key, tensor in _small_complete_top_level_tensors().items(): + _install_parameter(model, key, tensor) + _install_parameter(model, "x_embedder.junk", torch.ones(2)) + + with pytest.raises(ValueError, match="unexpected top-level.*x_embedder.junk"): + FluxSVDQuantNunchakuAdapter(require_complete_model=True).extra_tensors(model) + + +def test_complete_extra_collection_requires_and_copies_exact_top_level_parameters(): + model = _model({"num_layers": 0, "num_single_layers": 0}) + expected = _small_complete_top_level_tensors() + for key, tensor in expected.items(): + _install_parameter(model, key, tensor) + + actual = FluxSVDQuantNunchakuAdapter(require_complete_model=True).extra_tensors(model) + + assert set(actual) == set(expected) + assert all(tensor.dtype == torch.bfloat16 for tensor in actual.values()) + + +def test_complete_extra_collection_rejects_missing_top_level_parameter(): + model = _model({"num_layers": 0, "num_single_layers": 0}) + expected = _small_complete_top_level_tensors() + expected.pop("proj_out.bias") + for key, tensor in expected.items(): + _install_parameter(model, key, tensor) + + with pytest.raises(ValueError, match="missing top-level.*proj_out.bias"): + FluxSVDQuantNunchakuAdapter(require_complete_model=True).extra_tensors(model) + + +def test_generic_export_merges_adapter_extras_and_rejects_duplicates(): + residual = torch.nn.Linear(32, 8) + residual.data_type, residual.bits, residual.group_size, residual.sym = "mx_fp4", 4, 32, True + residual.act_data_type, residual.act_bits, residual.act_group_size = "mx_fp4", 4, 32 + residual.act_sym, residual.act_dynamic = True, True + from auto_round.algorithms.transforms.svdquant.wrapper import SVDQuantLinear + + wrapped = SVDQuantLinear( + residual, torch.nn.Linear(32, 2, bias=False), torch.nn.Linear(2, 8, bias=False), torch.ones(32) + ) + model = torch.nn.Sequential(wrapped) + + class ExtraAdapter(IdentitySVDQuantModelAdapter): + def extra_tensors(self, model): + return {"passthrough.weight": torch.ones(3).tanh()} + + tensors = collect_svdquant_tensors(model, adapter=ExtraAdapter()) + assert torch.equal(tensors["passthrough.weight"], torch.ones(3).tanh()) + + class DuplicateAdapter(IdentitySVDQuantModelAdapter): + def extra_tensors(self, model): + return {"0.bias": torch.ones(8)} + + with pytest.raises(ValueError, match="duplicate tensor key"): + collect_svdquant_tensors(model, adapter=DuplicateAdapter()) + + +def test_adapter_sources_have_no_external_runtime_imports(): + import auto_round.export.svdquant_adapters.flux as module + + source = inspect.getsource(module).lower() + assert "import deepcompressor" not in source + assert "import nunchaku" not in source + + +def test_partial_flux_collect_and_save_roundtrip(tmp_path): + from safetensors import safe_open + + from auto_round.algorithms.transforms.svdquant.wrapper import SVDQuantLinear + + residual = torch.nn.Linear(32, 8) + residual.data_type, residual.bits, residual.group_size, residual.sym = "mx_fp4", 4, 32, True + residual.act_data_type, residual.act_bits, residual.act_group_size = "mx_fp4", 4, 32 + residual.act_sym, residual.act_dynamic = True, True + wrapped = SVDQuantLinear( + residual, + torch.nn.Linear(32, 2, bias=False), + torch.nn.Linear(2, 8, bias=False), + torch.linspace(0.5, 1.5, 32), + ) + model = _model({"num_layers": 1, "num_single_layers": 0}) + _install(model, "transformer_blocks.0.attn.to_out.0", wrapped) + _install(model, "x_embedder", torch.nn.Linear(8, 8)) + adapter = FluxSVDQuantNunchakuAdapter(require_complete_model=False) + config = SVDQuantExportConfig(runtime_loadable=True) + + collected = collect_svdquant_tensors(model, adapter=adapter, config=config) + path = tmp_path / "flux.safetensors" + save_svdquant_nunchaku_safetensors(model, str(path), adapter=adapter, config=config) + + with safe_open(path, framework="pt") as handle: + assert set(handle.keys()) == set(collected) + assert handle.metadata()["model_class"] == "NunchakuFluxTransformer2dModel" + assert handle.get_tensor("x_embedder.weight").dtype == torch.bfloat16 diff --git a/test/test_cpu/export/test_svdquant_nunchaku_export.py b/test/test_cpu/export/test_svdquant_nunchaku_export.py new file mode 100644 index 0000000000..88deb2ed6a --- /dev/null +++ b/test/test_cpu/export/test_svdquant_nunchaku_export.py @@ -0,0 +1,590 @@ +import ast +import json +from dataclasses import replace +from pathlib import Path + +import pytest +import torch +from safetensors import safe_open + +from auto_round.algorithms.transforms.svdquant.wrapper import SVDQuantLinear +from auto_round.export.svdquant_mxfp4 import unpack_lowrank_weight +from auto_round.export.svdquant_nunchaku import ( + IdentitySVDQuantModelAdapter, + MXFP4ResidualTensorProvider, + SourceLinearRecord, + SVDQuantExportConfig, + SVDQuantExportRecord, + collect_svdquant_tensors, + pack_nunchaku_16bit_vector, + save_svdquant_nunchaku_safetensors, + unpack_nunchaku_16bit_vector, +) +from auto_round.wrapper import WrapperLinear + + +def _toy_model(*, rank=3, bias=True, in_features=65, out_features=7): + residual = torch.nn.Linear(in_features, out_features, bias=bias) + residual.data_type = "mx_fp4e2m1" + residual.bits = 4 + residual.group_size = 32 + residual.sym = True + residual.act_data_type = "mx_fp4e2m1" + residual.act_bits = 4 + residual.act_group_size = 32 + residual.act_sym = True + residual.act_dynamic = True + lora_down = torch.nn.Linear(in_features, rank, bias=False) + lora_up = torch.nn.Linear(rank, out_features, bias=False) + smooth = torch.arange(1, in_features + 1, dtype=torch.float32) + layer = SVDQuantLinear(residual, lora_down, lora_up, smooth) + return torch.nn.Sequential(layer) + + +def test_default_collection_emits_runtime_layout_tensors_without_debug_residual(): + model = _toy_model() + + tensors = collect_svdquant_tensors(model) + + assert set(tensors) == { + "0.qweight", + "0.wscales", + "0.smooth", + "0.smooth_orig", + "0.lora_down", + "0.lora_up", + "0.bias", + } + assert tensors["0.qweight"].shape == (128, 64) + assert tensors["0.qweight"].dtype == torch.int8 + assert tensors["0.wscales"].shape == (4, 128) + assert tensors["0.wscales"].dtype == torch.uint8 + for key in ("0.smooth", "0.smooth_orig", "0.lora_down", "0.lora_up", "0.bias"): + assert tensors[key].dtype == torch.bfloat16 + assert tensors["0.lora_down"].shape == (128, 16) + assert tensors["0.lora_up"].shape == (128, 16) + logical_down = model[0].lora_down.weight.detach().to(torch.bfloat16) + logical_up = model[0].lora_up.weight.detach().to(torch.bfloat16) + unpacked_down = unpack_lowrank_weight(tensors["0.lora_down"], down=True) + unpacked_up = unpack_lowrank_weight(tensors["0.lora_up"], down=False) + expected_runtime_down = logical_down * model[0].smooth.to(torch.bfloat16).unsqueeze(0) + torch.testing.assert_close(unpacked_down[:3, :65], expected_runtime_down) + torch.testing.assert_close(unpacked_up[:7, :3], logical_up) + torch.testing.assert_close( + unpack_nunchaku_16bit_vector(tensors["0.smooth"])[:65], model[0].smooth.reciprocal().to(torch.bfloat16) + ) + torch.testing.assert_close( + unpack_nunchaku_16bit_vector(tensors["0.bias"])[:7], + model[0].residual_linear.bias.detach().to(torch.bfloat16), + ) + + +def test_collection_reads_scheme_through_activation_quant_wrapper(): + model = _toy_model() + residual = model[0].residual_linear + residual.scale_dtype = torch.float16 + wrapped = WrapperLinear( + residual, + device="cpu", + enable_minmax_tuning=False, + enable_norm_bias_tuning=False, + enable_round_tuning=False, + disable_opt_rtn=True, + iters=0, + ).unwrapper({}) + model[0].residual_linear = wrapped + + assert not hasattr(wrapped, "bits") + assert wrapped.orig_layer.bits == 4 + + tensors = collect_svdquant_tensors(model) + + assert tensors["0.qweight"].dtype == torch.int8 + assert tensors["0.wscales"].dtype == torch.uint8 + + +def test_adapter_maps_all_logical_source_records_with_model_level_visibility(): + model = torch.nn.Sequential(_toy_model()[0], _toy_model()[0]) + + class CapturingAdapter(IdentitySVDQuantModelAdapter): + def map_modules(self, received_model, records): + assert received_model is model + records = tuple(records) + assert all(isinstance(record, SourceLinearRecord) for record in records) + assert [record.name for record in records] == ["0", "1"] + assert [record.lora_down.shape for record in records] == [(3, 65), (3, 65)] + assert [record.scheme.data_type for record in records] == ["mx_fp4e2m1", "mx_fp4e2m1"] + assert all(record.scheme.group_size == 32 and record.scheme.sym for record in records) + assert all(record.scheme.act_data_type == "mx_fp4e2m1" for record in records) + self.records = records + return super().map_modules(received_model, records) + + def validate_records(self, sources, records): + self.validated_records = (sources, records) + + adapter = CapturingAdapter() + + tensors = collect_svdquant_tensors(model, adapter=adapter) + + assert set(key.split(".", 1)[0] for key in tensors) == {"0", "1"} + assert len(adapter.records) == 2 + assert adapter.validated_records[0] is adapter.records + assert [record.prefix for record in adapter.validated_records[1]] == ["0", "1"] + + +def test_bias_and_smooth_vectors_match_layout_fixture_and_identity_padding(): + model = _toy_model(bias=False) + model[0].register_buffer("smooth_orig", torch.arange(1, 66, dtype=torch.float32)) + + tensors = collect_svdquant_tensors(model) + + bias = unpack_nunchaku_16bit_vector(tensors["0.bias"]) + assert torch.count_nonzero(bias[:7]) == 0 + assert torch.equal(bias[7:], torch.ones_like(bias[7:])) + smooth_orig = unpack_nunchaku_16bit_vector(tensors["0.smooth_orig"]) + torch.testing.assert_close( + smooth_orig[:65], torch.arange(1, 66, dtype=torch.float32).reciprocal().to(torch.bfloat16) + ) + assert torch.equal(smooth_orig[65:], torch.ones_like(smooth_orig[65:])) + fixture = pack_nunchaku_16bit_vector(torch.arange(128, dtype=torch.float16)) + assert fixture[:32].tolist() == [ + 0.0, + 1.0, + 8.0, + 9.0, + 2.0, + 3.0, + 10.0, + 11.0, + 4.0, + 5.0, + 12.0, + 13.0, + 6.0, + 7.0, + 14.0, + 15.0, + 16.0, + 17.0, + 24.0, + 25.0, + 18.0, + 19.0, + 26.0, + 27.0, + 20.0, + 21.0, + 28.0, + 29.0, + 22.0, + 23.0, + 30.0, + 31.0, + ] + + +def test_debug_unpacked_is_explicit_and_not_runtime_loadable(tmp_path): + config = SVDQuantExportConfig(debug_unpacked=True) + + tensors = collect_svdquant_tensors(_toy_model(), config=config) + + assert "0.residual.weight" in tensors + output_path = tmp_path / "invalid.safetensors" + with pytest.raises(ValueError, match="not runtime-loadable"): + save_svdquant_nunchaku_safetensors( + _toy_model(), output_path, config=SVDQuantExportConfig(debug_unpacked=True, runtime_loadable=True) + ) + assert not output_path.exists() + + +@pytest.mark.parametrize( + "field,value,message", + [ + ("weight_dtype", "mx_fp4e2m1", "weight_dtype"), + ("activation_dtype", "fp16", "activation_dtype"), + ("scale_dtype", "fp16", "scale_dtype"), + ("group_size", 16, "group_size"), + ("group_size", 32.0, "group_size"), + ("low_rank_dtype", torch.float32, "low_rank_dtype"), + ("debug_unpacked", 1, "debug_unpacked"), + ], +) +def test_config_rejects_non_nunchaku_formats(field, value, message): + with pytest.raises(ValueError, match=message): + SVDQuantExportConfig(**{field: value}) + + +@pytest.mark.parametrize("group_size", [32.0, True]) +def test_mxfp4_residual_provider_requires_non_bool_integer_group_size(group_size): + with pytest.raises(ValueError, match="group_size must be 32"): + MXFP4ResidualTensorProvider(group_size=group_size) + + +@pytest.mark.parametrize( + "field,value,message", + [ + ("data_type", "int", "data_type"), + ("bits", 8, "bits=4"), + ("bits", 4.0, "bits=4"), + ("group_size", (32, 32), "scalar group_size=32"), + ("group_size", 64, "scalar group_size=32"), + ("group_size", 32.0, "scalar group_size=32"), + ("sym", False, "sym=True"), + ("act_data_type", "int", "activation data_type"), + ("act_bits", 16, "activation bits=4"), + ("act_bits", 4.0, "activation bits=4"), + ("act_group_size", 64, "activation scalar group_size=32"), + ("act_group_size", 32.0, "activation scalar group_size=32"), + ("act_sym", False, "activation sym=True"), + ("act_dynamic", False, "act_dynamic=True"), + ], +) +def test_collection_rejects_incompatible_selected_scheme(field, value, message): + model = _toy_model() + setattr(model[0].residual_linear, field, value) + + with pytest.raises(ValueError, match=message): + collect_svdquant_tensors(model) + + +def test_collection_rejects_missing_selected_weight_or_activation_scheme(): + missing_weight = _toy_model() + del missing_weight[0].residual_linear.data_type + with pytest.raises(ValueError, match="missing.*data_type"): + collect_svdquant_tensors(missing_weight) + + missing_activation = _toy_model() + for field in ("act_data_type", "act_bits", "act_group_size", "act_sym", "act_dynamic"): + delattr(missing_activation[0].residual_linear, field) + with pytest.raises(ValueError, match="missing.*activation"): + collect_svdquant_tensors(missing_activation) + + +def test_collection_accepts_normal_autoround_mxfp4_preset_values(): + model = _toy_model() + model[0].residual_linear.data_type = "mx_fp" + model[0].residual_linear.act_data_type = "mx_fp" + + tensors = collect_svdquant_tensors(model) + + assert "0.qweight" in tensors + + +def test_collection_accepts_svdquant_rceil_deployment_types(): + model = _toy_model() + model[0].residual_linear.data_type = "mx_fp4e2m1_rceil" + model[0].residual_linear.act_data_type = "mx_fp4e2m1_rceil" + + tensors = collect_svdquant_tensors(model) + + assert "0.qweight" in tensors + + +def test_collection_rejects_nonfinite_values_and_mixed_ranks(): + nonfinite = _toy_model() + nonfinite[0].smooth[0] = torch.nan + with pytest.raises(ValueError, match="finite"): + collect_svdquant_tensors(nonfinite) + + mixed = torch.nn.Sequential(_toy_model(rank=2)[0], _toy_model(rank=3)[0]) + with pytest.raises(ValueError, match="mixed SVDQuant ranks"): + collect_svdquant_tensors(mixed) + + +def test_save_rejects_empty_model_and_missing_runtime_metadata_before_writing(tmp_path): + empty_path = tmp_path / "empty.safetensors" + with pytest.raises(ValueError, match="No SVDQuantLinear"): + save_svdquant_nunchaku_safetensors(torch.nn.Linear(2, 2), empty_path) + assert not empty_path.exists() + + runtime_path = tmp_path / "runtime.safetensors" + with pytest.raises(ValueError, match="model_class.*serialized 'config'"): + save_svdquant_nunchaku_safetensors( + _toy_model(), runtime_path, config=SVDQuantExportConfig(runtime_loadable=True) + ) + assert not runtime_path.exists() + + +def test_runtime_metadata_validation_precedes_source_cpu_copy_and_packing(monkeypatch, tmp_path): + class Provider: + def tensors_for(self, record): + pytest.fail("residual packing must not run before runtime metadata validation") + + def reject_cpu_copy(self, *args, **kwargs): + pytest.fail("source tensors must remain on their original device before serialization") + + monkeypatch.setattr(torch.Tensor, "cpu", reject_cpu_copy) + + with pytest.raises(ValueError, match="model_class.*serialized 'config'"): + save_svdquant_nunchaku_safetensors( + _toy_model(), + tmp_path / "runtime.safetensors", + config=SVDQuantExportConfig(runtime_loadable=True), + residual_provider=Provider(), + ) + + +def test_collection_retains_source_device_storage_until_after_packing(monkeypatch): + model = _toy_model() + events = [] + original_cpu = torch.Tensor.cpu + + class Adapter(IdentitySVDQuantModelAdapter): + def map_modules(self, received_model, records): + records = tuple(records) + events.append("map") + assert records[0].residual_weight.device == model[0].residual_linear.weight.device + assert records[0].residual_weight.data_ptr() == model[0].residual_linear.weight.data_ptr() + return super().map_modules(received_model, records) + + class Provider: + def __init__(self): + self.delegate = MXFP4ResidualTensorProvider() + + def tensors_for(self, record): + events.append("pack") + return self.delegate.tensors_for(record) + + def tracked_cpu(self, *args, **kwargs): + events.append("cpu") + return original_cpu(self, *args, **kwargs) + + monkeypatch.setattr(torch.Tensor, "cpu", tracked_cpu) + + collect_svdquant_tensors(model, adapter=Adapter(), residual_provider=Provider()) + + assert events[0] == "map" + assert events.index("pack") < events.index("cpu") + + +def test_runtime_adapter_metadata_and_validation_receive_resolved_rank(tmp_path): + class RuntimeAdapter(IdentitySVDQuantModelAdapter): + def metadata(self, model, rank): + self.metadata_rank = rank + return {"model_class": "ToyModel", "config": json.dumps({"hidden_size": 65})} + + def validate(self, tensors, metadata): + self.validated = (set(tensors), metadata["model_class"]) + + adapter = RuntimeAdapter() + output_path = tmp_path / "runtime.safetensors" + + save_svdquant_nunchaku_safetensors( + _toy_model(), output_path, config=SVDQuantExportConfig(runtime_loadable=True), adapter=adapter + ) + + assert adapter.metadata_rank == 3 + assert adapter.validated[1] == "ToyModel" + assert "0.qweight" in adapter.validated[0] + + +def test_adapter_can_expand_records_and_remap_export_suffixes(): + class ExpandingAdapter(IdentitySVDQuantModelAdapter): + def map_modules(self, model, records): + source = tuple(records)[0] + values = { + "residual_weight": source.residual_weight, + "lora_down": source.lora_down, + "lora_up": source.lora_up, + "smooth": source.smooth, + "smooth_orig": source.smooth_orig, + "bias": source.bias, + "scheme": source.scheme, + "sources": (source,), + } + return ( + SVDQuantExportRecord(prefix="left", key_mapping={"qweight": "packed_weight"}, **values), + SVDQuantExportRecord(prefix="right", **values), + ) + + tensors = collect_svdquant_tensors(_toy_model(), adapter=ExpandingAdapter()) + + assert "left.packed_weight" in tensors + assert "left.qweight" not in tensors + assert "right.qweight" in tensors + + +def test_adapter_can_recompose_independent_sources_at_configured_rank(): + model = torch.nn.Sequential(_toy_model()[0], _toy_model()[0]) + model[0].residual_linear.data_type = "mx_fp4" + model[0].residual_linear.act_data_type = "mx_fp4" + + assert not torch.equal(model[0].lora_down.weight, model[1].lora_down.weight) + assert not torch.equal(model[0].lora_up.weight, model[1].lora_up.weight) + + class FusionAdapter(IdentitySVDQuantModelAdapter): + fused_record = None + + def map_modules(self, model, records): + left, right = tuple(records) + rank = left.lora_down.shape[0] + effective_weights = tuple( + source.residual_weight + source.lora_up @ source.lora_down for source in (left, right) + ) + fused_weight = torch.cat(effective_weights, dim=0) + u, singular_values, vh = torch.linalg.svd(fused_weight.float(), full_matrices=False) + lora_up = (u[:, :rank] * singular_values[:rank]).to(fused_weight.dtype) + lora_down = vh[:rank].to(fused_weight.dtype) + self.fused_record = SVDQuantExportRecord( + prefix="fused", + residual_weight=fused_weight - lora_up @ lora_down, + lora_down=lora_down, + lora_up=lora_up, + smooth=left.smooth, + smooth_orig=left.smooth_orig, + bias=torch.cat((left.bias, right.bias), dim=0), + scheme=left.scheme, + sources=(left, right), + ) + return (self.fused_record,) + + adapter = FusionAdapter() + tensors = collect_svdquant_tensors(model, adapter=adapter) + + assert "fused.qweight" in tensors + assert tensors["fused.qweight"].shape == (128, 64) + assert adapter.fused_record.lora_down.shape[0] == 3 + expected = torch.cat( + tuple( + module.residual_linear.weight.detach() + module.lora_up.weight.detach() @ module.lora_down.weight.detach() + for module in model + ), + dim=0, + ) + actual = adapter.fused_record.residual_weight + adapter.fused_record.lora_up @ adapter.fused_record.lora_down + torch.testing.assert_close(actual, expected) + + +def test_adapter_provenance_rejects_dropped_and_foreign_sources(): + model = torch.nn.Sequential(_toy_model()[0], _toy_model()[0]) + + class DroppingAdapter(IdentitySVDQuantModelAdapter): + def map_modules(self, model, records): + return super().map_modules(model, tuple(records)[:1]) + + with pytest.raises(ValueError, match="dropped logical sources.*1"): + collect_svdquant_tensors(model, adapter=DroppingAdapter()) + + class ForeignAdapter(IdentitySVDQuantModelAdapter): + def map_modules(self, model, records): + source = tuple(records)[0] + record = tuple(super().map_modules(model, (source,)))[0] + return (replace(record, sources=(replace(source, name="foreign"),)),) + + with pytest.raises(ValueError, match="foreign logical source"): + collect_svdquant_tensors(_toy_model(), adapter=ForeignAdapter()) + + +def test_adapter_provenance_requires_source_rank_to_equal_output_rank(): + class RankChangingAdapter(IdentitySVDQuantModelAdapter): + def map_modules(self, model, records): + source = tuple(records)[0] + record = tuple(super().map_modules(model, (source,)))[0] + return ( + replace( + record, + lora_down=record.lora_down[:2], + lora_up=record.lora_up[:, :2], + ), + ) + + with pytest.raises(ValueError, match="Nunchaku.*configured rank.*exact rank-sum fusion is unsupported"): + collect_svdquant_tensors(_toy_model(), adapter=RankChangingAdapter()) + + +def test_identity_adapter_uses_stable_model_prefix_for_root_svdquant_linear(): + root = _toy_model()[0] + + tensors = collect_svdquant_tensors(root) + + assert "model.qweight" in tensors + assert not any(key.startswith(".") for key in tensors) + + +@pytest.mark.parametrize( + "payload,message", + [ + ({"qweight": torch.zeros(128, 64), "wscales": torch.zeros(4, 128, dtype=torch.uint8)}, "qweight"), + ({"qweight": torch.zeros(128, 64, dtype=torch.int8)}, "qweight.*wscales"), + ( + { + "qweight": torch.zeros(128, 64, dtype=torch.int8), + "wscales": torch.zeros(3, 128, dtype=torch.uint8), + }, + "wscales shape", + ), + ], +) +def test_collection_rejects_malformed_packed_residual_payloads(payload, message): + class Provider: + def tensors_for(self, record): + return payload + + with pytest.raises(ValueError, match=message): + collect_svdquant_tensors(_toy_model(), residual_provider=Provider()) + + +def test_collection_rejects_aligned_payload_that_is_too_small_for_logical_record(): + class WrongProvider: + def tensors_for(self, record): + return { + "qweight": torch.zeros(128, 64, dtype=torch.int8), + "wscales": torch.zeros(4, 128, dtype=torch.uint8), + } + + with pytest.raises(ValueError, match=r"qweight shape.*\(256, 128\)"): + collect_svdquant_tensors(_toy_model(in_features=129, out_features=129), residual_provider=WrongProvider()) + + +def test_save_svdquant_nunchaku_safetensors_writes_metadata(tmp_path): + output_path = tmp_path / "svdquant.safetensors" + + save_svdquant_nunchaku_safetensors(_toy_model(), str(output_path)) + + with safe_open(output_path, framework="pt") as handle: + keys = set(handle.keys()) + metadata = handle.metadata() + quantization_config = json.loads(metadata["quantization_config"]) + + assert "0.qweight" in keys + assert "0.wscales" in keys + assert "0.residual.weight" not in keys + assert metadata == { + "artifact_type": "generic_intermediate", + "quantization_config": json.dumps( + { + "method": "svdquant", + "weight": {"dtype": "fp4_e2m1_all", "scale_dtype": "ue8m0", "group_size": 32}, + "activation": {"dtype": "fp4_e2m1_all", "scale_dtype": "ue8m0", "group_size": 32}, + "rank": 3, + }, + sort_keys=True, + ), + } + assert quantization_config == { + "method": "svdquant", + "weight": {"dtype": "fp4_e2m1_all", "scale_dtype": "ue8m0", "group_size": 32}, + "activation": {"dtype": "fp4_e2m1_all", "scale_dtype": "ue8m0", "group_size": 32}, + "rank": 3, + } + + +def test_svdquant_nunchaku_exporter_has_no_runtime_project_imports(): + import auto_round + + forbidden = {"deepcompressor", "nunchaku"} + violations = [] + source_root = Path(auto_round.__file__).parent + for path in source_root.rglob("*.py"): + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + for node in ast.walk(tree): + if isinstance(node, ast.Import): + names = {alias.name.split(".", 1)[0] for alias in node.names} + elif isinstance(node, ast.ImportFrom) and node.module: + names = {node.module.split(".", 1)[0]} + else: + continue + if names & forbidden: + violations.append(f"{path.relative_to(source_root)}:{node.lineno}") + + assert violations == [] diff --git a/test/test_cpu/export/test_svdquant_nunchaku_format.py b/test/test_cpu/export/test_svdquant_nunchaku_format.py new file mode 100644 index 0000000000..3ff45fc175 --- /dev/null +++ b/test/test_cpu/export/test_svdquant_nunchaku_format.py @@ -0,0 +1,241 @@ +# Copyright (c) 2026 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +from safetensors.torch import save_file + +from auto_round.algorithms.transforms.svdquant.wrapper import SVDQuantLinear +from auto_round.compressors.base import BaseCompressor +from auto_round.formats import SVDQuantNunchakuFormat, get_formats +from auto_round.schemes import PRESET_SCHEMES + + +def _mxfp4_compressor(**updates): + values = PRESET_SCHEMES["MXFP4"].to_dict() + values.update(scheme="MXFP4", **updates) + return SimpleNamespace(**values) + + +def _toy_svd_model(): + model = torch.nn.Module() + residual = torch.nn.Linear(32, 32) + residual.data_type = "mx_fp4e2m1" + residual.bits = 4 + residual.group_size = 32 + residual.sym = True + residual.act_data_type = "mx_fp4e2m1" + residual.act_bits = 4 + residual.act_group_size = 32 + residual.act_sym = True + residual.act_dynamic = True + model.svd = SVDQuantLinear( + residual, + torch.nn.Linear(32, 1, bias=False), + torch.nn.Linear(1, 32, bias=False), + torch.ones(32), + ) + return model + + +def test_get_formats_resolves_full_model_svdquant_nunchaku_format(): + output_format = get_formats("svdquant_nunchaku", _mxfp4_compressor())[0] + + assert output_format.format_name == "svdquant_nunchaku" + assert output_format.requires_full_model_export is True + + +def test_svdquant_nunchaku_rejects_incompatible_scheme(): + scheme = PRESET_SCHEMES["MXFP4"].copy() + scheme.group_size = 64 + + with pytest.raises(ValueError, match=r"group_size=64.*group_size=32"): + SVDQuantNunchakuFormat.check_scheme_args(scheme) + + +def test_full_model_format_disables_immediate_packing_and_saving(): + output_format = get_formats("svdquant_nunchaku", _mxfp4_compressor())[0] + compressor = SimpleNamespace( + formats=[output_format], + inplace=True, + has_qlayer_outside_block=False, + need_calib=True, + model_context=SimpleNamespace(model=torch.nn.Module(), is_mllm=False), + compress_context=SimpleNamespace( + low_cpu_mem_usage=True, + is_immediate_packing=True, + is_immediate_saving=True, + ), + quantize_config=SimpleNamespace(data_type="mx_fp"), + output_dir="unused", + _ensure_shard_writer=lambda: pytest.fail("full-model export must not create a shard writer"), + ) + + BaseCompressor._adjust_immediate_packing_and_saving(compressor) + + assert compressor.compress_context.is_immediate_packing is False + assert compressor.compress_context.is_immediate_saving is False + + +def test_format_uses_flux_adapter_and_diffusers_weight_name(monkeypatch, tmp_path): + import auto_round.export.svdquant_nunchaku as exporter + from auto_round.export.svdquant_adapters.flux import FluxSVDQuantNunchakuAdapter + + output_format = get_formats("svdquant_nunchaku", _mxfp4_compressor())[0] + model = torch.nn.Module() + model.config = {"_class_name": "FluxTransformer2DModel", "num_layers": 0, "num_single_layers": 0} + captured = {} + + def fake_export(export_model, output_path, *, config, residual_provider, adapter): + captured.update( + model=export_model, + output_path=output_path, + runtime_loadable=config.runtime_loadable, + residual_provider=residual_provider, + adapter=adapter, + ) + + monkeypatch.setattr(exporter, "save_svdquant_nunchaku_safetensors", fake_export) + + result = output_format.save_quantized(tmp_path, model=model, model_adapter="flux", device="cpu") + + assert result is model + assert captured["output_path"] == str(tmp_path / "diffusion_pytorch_model.safetensors") + assert captured["runtime_loadable"] is True + assert isinstance(captured["adapter"], FluxSVDQuantNunchakuAdapter) + + +def test_format_rejects_models_without_runtime_adapter(tmp_path): + output_format = get_formats("svdquant_nunchaku", _mxfp4_compressor())[0] + + with pytest.raises(ValueError, match="runtime model adapter"): + output_format.save_quantized(tmp_path, model=torch.nn.Linear(2, 2), model_adapter="auto") + + +def test_format_rejects_incompatible_residual_override(monkeypatch, tmp_path): + import auto_round.export.svdquant_nunchaku as exporter + + output_format = get_formats("svdquant_nunchaku", _mxfp4_compressor())[0] + monkeypatch.setattr( + exporter, + "save_svdquant_nunchaku_safetensors", + lambda *args, **kwargs: pytest.fail("exporter must not be called"), + ) + + with pytest.raises(ValueError, match=r"group_size=64.*group_size=32"): + output_format.save_quantized( + tmp_path, + model=_toy_svd_model(), + layer_config={"svd.residual_linear": {"group_size": 64}}, + ) + + +def test_diffusion_save_exports_self_contained_nunchaku_pipeline(tmp_path): + from auto_round.compressors.diffusion_mixin import DiffusionMixin + + class QuantizedTransformer(torch.nn.Module): + def save_config(self, output_dir): + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + (output_dir / "config.json").write_text('{"_class_name": "FluxTransformer2DModel"}', encoding="utf-8") + + class Bf16Component: + def save_pretrained(self, output_dir): + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + (output_dir / "config.json").write_text("{}", encoding="utf-8") + (output_dir / "diffusion_pytorch_model.safetensors").touch() + + transformer = QuantizedTransformer() + vae = Bf16Component() + + class Pipeline: + def __init__(self): + self.transformer = transformer + self.vae = vae + self.components = {"transformer": transformer, "vae": vae} + + def save_config(self, output_dir): + Path(output_dir, "model_index.json").write_text( + json.dumps( + { + "_class_name": "FluxPipeline", + "transformer": ["diffusers", "FluxTransformer2DModel"], + "vae": ["diffusers", "AutoencoderKL"], + } + ), + encoding="utf-8", + ) + + class ExportParent: + def save_quantized(self, output_dir, **kwargs): + Path(output_dir).mkdir(parents=True, exist_ok=True) + save_file( + {"probe": torch.zeros(1)}, + f"{output_dir}/diffusion_pytorch_model.safetensors", + metadata={"model_class": "NunchakuFluxTransformer2dModel"}, + ) + return self.model_context.model + + class Compressor(DiffusionMixin, ExportParent): + pass + + compressor = Compressor.__new__(Compressor) + compressor.formats = [SimpleNamespace(format_name="svdquant_nunchaku")] + compressor.model_context = SimpleNamespace(pipe=Pipeline(), model=transformer) + compressor.compress_context = SimpleNamespace(is_immediate_saving=False) + + compressor.save_quantized(tmp_path) + + assert (tmp_path / "transformer" / "diffusion_pytorch_model.safetensors").is_file() + assert (tmp_path / "vae" / "diffusion_pytorch_model.safetensors").is_file() + model_index = json.loads((tmp_path / "model_index.json").read_text(encoding="utf-8")) + assert model_index["transformer"] == ["nunchaku", "NunchakuFluxTransformer2dModel"] + assert model_index["vae"] == ["diffusers", "AutoencoderKL"] + + +def test_diffusion_save_requires_runtime_model_class_metadata(tmp_path): + from auto_round.compressors.diffusion_mixin import DiffusionMixin + + model = torch.nn.Module() + + class Pipeline: + transformer = model + components = {"transformer": model} + + def save_config(self, output_dir): + Path(output_dir, "model_index.json").write_text( + json.dumps({"transformer": ["diffusers", "FluxTransformer2DModel"]}), encoding="utf-8" + ) + + class ExportParent: + def save_quantized(self, output_dir, **kwargs): + Path(output_dir).mkdir(parents=True, exist_ok=True) + save_file({"probe": torch.zeros(1)}, f"{output_dir}/diffusion_pytorch_model.safetensors") + return model + + class Compressor(DiffusionMixin, ExportParent): + pass + + compressor = Compressor.__new__(Compressor) + compressor.formats = [SimpleNamespace(format_name="svdquant_nunchaku")] + compressor.model_context = SimpleNamespace(pipe=Pipeline(), model=model) + compressor.compress_context = SimpleNamespace(is_immediate_saving=False) + + with pytest.raises(ValueError, match="model_class"): + compressor.save_quantized(tmp_path) From f0c63d9c2290247bbc807719ef153829c69e5a3f Mon Sep 17 00:00:00 2001 From: changwangss Date: Tue, 28 Jul 2026 17:52:06 +0800 Subject: [PATCH 7/9] fix: support low-memory diffusion SignRound calibration Signed-off-by: changwangss --- auto_round/algorithms/block_runner.py | 8 --- .../algorithms/transforms/svdquant/apply.py | 22 ++++++- auto_round/calibration/diffusion.py | 23 ++++++- auto_round/compressors/diffusion_mixin.py | 4 ++ auto_round/formats.py | 9 +++ auto_round/utils/model.py | 36 +++++------ test/test_cpu/algorithms/test_block_runner.py | 62 +++++++++++++++++++ test/test_cpu/algorithms/test_svdquant.py | 28 +++++++++ .../export/test_svdquant_nunchaku_format.py | 12 +++- test/test_cpu/models/test_diffusion.py | 40 ++++++++++++ .../utils/test_diffusion_detection.py | 26 ++++++++ 11 files changed, 236 insertions(+), 34 deletions(-) create mode 100644 test/test_cpu/algorithms/test_block_runner.py create mode 100644 test/test_cpu/utils/test_diffusion_detection.py diff --git a/auto_round/algorithms/block_runner.py b/auto_round/algorithms/block_runner.py index 0e88ff3b42..5eebde7f66 100644 --- a/auto_round/algorithms/block_runner.py +++ b/auto_round/algorithms/block_runner.py @@ -232,14 +232,6 @@ def forward( self.last_output_dict["hidden_states"] = result return result else: - if self.batch_size == 1: - outputs = [output.unsqueeze(dim=self.batch_dim).to(out_device) for output in outputs] - if output_dict: - output_dict = { - key: [value.unsqueeze(dim=self.batch_dim).to(out_device) for value in values] - for key, values in output_dict.items() - } - outputs = torch.cat(outputs, dim=self.batch_dim).to(out_device) if output_dict: self.last_output_dict = { diff --git a/auto_round/algorithms/transforms/svdquant/apply.py b/auto_round/algorithms/transforms/svdquant/apply.py index 539ba61c16..9fb7391002 100644 --- a/auto_round/algorithms/transforms/svdquant/apply.py +++ b/auto_round/algorithms/transforms/svdquant/apply.py @@ -143,6 +143,7 @@ def __init__(self, config: SVDQuantConfig) -> None: self._configured_block_names: tuple[str, ...] = () self._block_groups: dict[str, list[SmoothSearchGroup]] = {} self._smooth_calibration: dict[str, SmoothGroupCalibration] = {} + self._target_modules = config.target_modules def bind(self, orchestrator) -> None: super().bind(orchestrator) @@ -158,7 +159,22 @@ def prepare_run(self, composer=None) -> None: self._block_groups.clear() if self.model is None: return - self.model._autoround_svdquant_model_adapter = self.config.model_adapter or "auto" + model_adapter = self.config.model_adapter or "auto" + if model_adapter == "auto": + model_config = getattr(self.model, "config", None) + class_name = ( + model_config.get("_class_name", type(self.model).__name__) + if hasattr(model_config, "get") + else type(self.model).__name__ + ) + if "fluxtransformer" in str(class_name).lower(): + model_adapter = "flux" + self.model._autoround_svdquant_model_adapter = model_adapter + self._target_modules = self.config.target_modules + if self._target_modules is None and model_adapter == "flux": + from auto_round.export.svdquant_adapters import FLUX_SVDQUANT_TARGET_MODULES + + self._target_modules = FLUX_SVDQUANT_TARGET_MODULES for block_name in self._configured_block_names: block = self.model.get_submodule(block_name) self._block_groups[block_name] = discover_svdquant_groups(block, self._is_target) @@ -631,8 +647,8 @@ def _is_target(self, name: str, module: torch.nn.Module) -> bool: if not isinstance(module, torch.nn.Linear): return False full_name = str(getattr(module, "global_name", name)) - if self.config.target_modules and not any( - pattern in name or pattern in full_name for pattern in self.config.target_modules + if self._target_modules and not any( + pattern in name or pattern in full_name for pattern in self._target_modules ): return False if self.config.exclude_modules and any( diff --git a/auto_round/calibration/diffusion.py b/auto_round/calibration/diffusion.py index 1ca5278f1a..817816d801 100644 --- a/auto_round/calibration/diffusion.py +++ b/auto_round/calibration/diffusion.py @@ -32,6 +32,22 @@ from auto_round.compressors import BaseOrchestrator as BaseCompressor from auto_round.logger import logger from auto_round.utils.device_manager import device_manager + + +def _prepare_pipeline_for_calibration(pipe, target_device, *, low_gpu_mem_usage: bool) -> str | None: + """Place a diffusion pipeline for calibration without exceeding one GPU.""" + target_device = torch.device(target_device) + if low_gpu_mem_usage: + enable_model_cpu_offload = getattr(pipe, "enable_model_cpu_offload", None) + if not callable(enable_model_cpu_offload): + raise ValueError("The diffusion pipeline does not support component-level model CPU offload.") + enable_model_cpu_offload(device=target_device) + return "model" + if pipe.device != target_device: + pipe.to(target_device) + return None + + from auto_round.utils.model import wrap_block_forward_positional_to_kwargs @@ -116,8 +132,11 @@ def calib(self, nsamples: int, bs: int) -> None: exit(-1) target_device = device_manager.device - if pipe.device != torch.device(target_device): - pipe.to(target_device) + self._cpu_offload_mode = _prepare_pipeline_for_calibration( + pipe, + target_device, + low_gpu_mem_usage=self.low_gpu_mem_usage, + ) pipeline_fn = getattr(pipe, "_autoround_pipeline_fn", None) # Check if this is an I2V pipeline (needs calibration image) requires_image = False diff --git a/auto_round/compressors/diffusion_mixin.py b/auto_round/compressors/diffusion_mixin.py index 8273a86843..4c80c3dfca 100644 --- a/auto_round/compressors/diffusion_mixin.py +++ b/auto_round/compressors/diffusion_mixin.py @@ -396,6 +396,10 @@ def quantize(self) -> tuple[torch.nn.Module, dict]: layer_names=[], ) self.inputs = all_inputs + if getattr(self.calibration, "_cpu_offload_mode", None) == "model": + from accelerate.hooks import remove_hook_from_submodules + + remove_hook_from_submodules(self.model_context.model) clear_memory() self._inputs_cached = True return super().quantize() diff --git a/auto_round/formats.py b/auto_round/formats.py index 02a42cd74a..c487afd7c4 100644 --- a/auto_round/formats.py +++ b/auto_round/formats.py @@ -451,6 +451,15 @@ def save_quantized( ) self._validate_svd_layer_overrides(model, layer_config) + save_config = getattr(model, "save_config", None) + if callable(save_config): + save_config(output_dir) + else: + model_config = getattr(model, "config", None) + save_pretrained = getattr(model_config, "save_pretrained", None) + if callable(save_pretrained): + save_pretrained(output_dir) + model_adapter = model_adapter or getattr(model, "_autoround_svdquant_model_adapter", "auto") if isinstance(model_adapter, str): from auto_round.export.svdquant_adapters import resolve_svdquant_model_adapter diff --git a/auto_round/utils/model.py b/auto_round/utils/model.py index d84b18b64a..cc5f48b0bf 100644 --- a/auto_round/utils/model.py +++ b/auto_round/utils/model.py @@ -1150,42 +1150,40 @@ def is_gguf_model(model_path: Union[str, torch.nn.Module]) -> bool: def is_diffusion_model(model_or_path: Union[str, object], trust_remote_code: bool = True) -> bool: from auto_round.utils.common import LazyImport - # Then check if model_index.json exists for diffusion pipeline, - # which is a strong signal of being a diffusion pipeline. if isinstance(model_or_path, str): - # Quick check to avoid config loading attempts and unnecessary warnings if is_gguf_model(model_or_path): return False - # First check if it's a known diffusion pipeline by config/model_type - # to avoid unnecessary imports and file checks for non-diffusion models, which can be time-consuming. + if os.path.isdir(model_or_path): + index_file = os.path.join(model_or_path, "model_index.json") + if os.path.isfile(index_file): + check_diffusers_installed() + return True + elif os.path.isabs(model_or_path) or model_or_path.startswith(("./", "../")): + # Missing local pipeline components are not Hugging Face repo ids. + return False + + # NextStep is a diffusion model without the standard pipeline index. try: from transformers import AutoConfig config = AutoConfig.from_pretrained(model_or_path, trust_remote_code=trust_remote_code) model_type = getattr(config, "model_type", "") - # A special case for NextStep if model_type == "nextstep": return True - except: - logger.warning( - f"Failed to load config for {model_or_path}, trying to check model_index.json for diffusion pipeline." - ) - index_file = None + except Exception: + logger.debug("Failed to load AutoConfig while checking whether %s is a diffusion model", model_or_path) + if not os.path.isdir(model_or_path): try: from huggingface_hub import hf_hub_download index_file = hf_hub_download(model_or_path, "model_index.json") check_diffusers_installed() - except Exception as e: - print(e) - index_file = None - - elif os.path.exists(os.path.join(model_or_path, "model_index.json")): - check_diffusers_installed() - index_file = os.path.join(model_or_path, "model_index.json") - return index_file is not None + return index_file is not None + except Exception: + logger.debug("No model_index.json found while checking %s", model_or_path) + return False elif not isinstance(model_or_path, torch.nn.Module): check_diffusers_installed() pipeline_utils = LazyImport("diffusers.pipelines.pipeline_utils") diff --git a/test/test_cpu/algorithms/test_block_runner.py b/test/test_cpu/algorithms/test_block_runner.py new file mode 100644 index 0000000000..d0e4cb1670 --- /dev/null +++ b/test/test_cpu/algorithms/test_block_runner.py @@ -0,0 +1,62 @@ +# Copyright (c) 2026 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +from auto_round.algorithms.block_runner import BlockForwardRunner + + +def test_indexed_single_sample_forward_preserves_one_batch_dimension(): + runner = BlockForwardRunner( + batch_dim=0, + batch_size=1, + device="cpu", + cache_device="cpu", + amp=True, + amp_dtype=torch.bfloat16, + ) + sample = torch.randn(1, 3, 4) + + output = runner(torch.nn.Identity(), [sample], {}, indices=torch.tensor([0])) + + assert output.shape == sample.shape + torch.testing.assert_close(output, sample) + + +def test_indexed_diffusion_outputs_preserve_batch_dimension(): + class FluxTransformerBlock(torch.nn.Module): + def forward(self, hidden_states, **_kwargs): + return hidden_states + 1, hidden_states + 2 + + runner = BlockForwardRunner( + batch_dim=0, + batch_size=1, + device="cpu", + cache_device="cpu", + amp=True, + amp_dtype=torch.bfloat16, + is_diffusion=True, + ) + sample = torch.randn(1, 3, 4) + + output = runner( + FluxTransformerBlock(), + {"hidden_states": [sample]}, + {}, + indices=torch.tensor([0]), + ) + + assert output.shape == sample.shape + assert runner.last_output_dict["encoder_hidden_states"].shape == sample.shape + assert runner.last_output_dict["hidden_states"].shape == sample.shape diff --git a/test/test_cpu/algorithms/test_svdquant.py b/test/test_cpu/algorithms/test_svdquant.py index 7b010a1408..dec14e4a95 100644 --- a/test/test_cpu/algorithms/test_svdquant.py +++ b/test/test_cpu/algorithms/test_svdquant.py @@ -34,6 +34,8 @@ class FluxTransformerBlock(torch.nn.Module): def __init__(self, width=8): super().__init__() self.attn = FluxAttention(width) + self.norm1 = torch.nn.Module() + self.norm1.linear = torch.nn.Linear(width, width, bias=False) class TinyFlux(torch.nn.Module): @@ -103,6 +105,32 @@ def test_no_smooth_flux_qkv_share_one_down_factor(): torch.testing.assert_close(actual, reference) +def test_flux_adapter_default_targets_only_runtime_supported_projections(): + model = TinyFlux() + model.config = {"_class_name": "FluxTransformer2DModel"} + _mark_modules(model) + block_name = "transformer_blocks.0" + transform = SVDQuantTransform(SVDQuantConfig(rank=2, model_adapter="auto", low_rank_dtype="fp32")) + orchestrator = SimpleNamespace( + model_context=SimpleNamespace(model=model), + compress_context=None, + calibration_context=None, + scheme_context=None, + scale_dtype=None, + nblocks=1, + quant_block_list=[[block_name]], + ) + transform.bind(orchestrator) + transform.prepare_run() + + transform.pre_quantize_block( + BlockContext(model=model, block_names=[block_name], block_name=block_name, block_index=0) + ) + + assert isinstance(model.transformer_blocks[0].attn.to_q, SVDQuantLinear) + assert isinstance(model.transformer_blocks[0].norm1.linear, torch.nn.Linear) + + def test_no_smooth_grouped_residual_iteration_and_cleanup(): model = TinyFlux(width=32) _mark_modules(model) diff --git a/test/test_cpu/export/test_svdquant_nunchaku_format.py b/test/test_cpu/export/test_svdquant_nunchaku_format.py index 3ff45fc175..553db01819 100644 --- a/test/test_cpu/export/test_svdquant_nunchaku_format.py +++ b/test/test_cpu/export/test_svdquant_nunchaku_format.py @@ -97,8 +97,15 @@ def test_format_uses_flux_adapter_and_diffusers_weight_name(monkeypatch, tmp_pat from auto_round.export.svdquant_adapters.flux import FluxSVDQuantNunchakuAdapter output_format = get_formats("svdquant_nunchaku", _mxfp4_compressor())[0] - model = torch.nn.Module() - model.config = {"_class_name": "FluxTransformer2DModel", "num_layers": 0, "num_single_layers": 0} + + class FluxModel(torch.nn.Module): + config = {"_class_name": "FluxTransformer2DModel", "num_layers": 0, "num_single_layers": 0} + + def save_config(self, output_dir): + Path(output_dir).mkdir(parents=True, exist_ok=True) + Path(output_dir, "config.json").write_text(json.dumps(self.config), encoding="utf-8") + + model = FluxModel() captured = {} def fake_export(export_model, output_path, *, config, residual_provider, adapter): @@ -118,6 +125,7 @@ def fake_export(export_model, output_path, *, config, residual_provider, adapter assert captured["output_path"] == str(tmp_path / "diffusion_pytorch_model.safetensors") assert captured["runtime_loadable"] is True assert isinstance(captured["adapter"], FluxSVDQuantNunchakuAdapter) + assert (tmp_path / "config.json").is_file() def test_format_rejects_models_without_runtime_adapter(tmp_path): diff --git a/test/test_cpu/models/test_diffusion.py b/test/test_cpu/models/test_diffusion.py index 3d047cb72a..d92575f7d3 100644 --- a/test/test_cpu/models/test_diffusion.py +++ b/test/test_cpu/models/test_diffusion.py @@ -6,12 +6,52 @@ from packaging import version from auto_round import AutoRound +from auto_round.calibration.diffusion import _prepare_pipeline_for_calibration from ...helpers import get_model_path, transformers_version flux_name_or_path = get_model_path("black-forest-labs/FLUX.1-dev") +def test_low_gpu_memory_diffusion_calibration_uses_model_cpu_offload(): + class Pipeline: + device = torch.device("cpu") + + def __init__(self): + self.offload_device = None + + def enable_model_cpu_offload(self, *, device): + self.offload_device = device + + def to(self, _device): + raise AssertionError("low GPU memory calibration must not move the full pipeline") + + pipe = Pipeline() + + mode = _prepare_pipeline_for_calibration(pipe, "cuda:0", low_gpu_mem_usage=True) + + assert mode == "model" + assert pipe.offload_device == torch.device("cuda:0") + + +def test_regular_diffusion_calibration_moves_pipeline_to_device(): + class Pipeline: + device = torch.device("cpu") + + def __init__(self): + self.target_device = None + + def to(self, device): + self.target_device = device + + pipe = Pipeline() + + mode = _prepare_pipeline_for_calibration(pipe, "cuda:0", low_gpu_mem_usage=False) + + assert mode is None + assert pipe.target_device == torch.device("cuda:0") + + @pytest.fixture def setup_flux(): """Fixture to set up the Flux model and tokenizer.""" diff --git a/test/test_cpu/utils/test_diffusion_detection.py b/test/test_cpu/utils/test_diffusion_detection.py new file mode 100644 index 0000000000..8d6804e5a9 --- /dev/null +++ b/test/test_cpu/utils/test_diffusion_detection.py @@ -0,0 +1,26 @@ +from pathlib import Path + +from auto_round.utils.model import is_diffusion_model + + +def test_local_diffusion_pipeline_is_detected_from_model_index(monkeypatch, tmp_path): + (tmp_path / "model_index.json").write_text("{}", encoding="utf-8") + monkeypatch.setattr("auto_round.utils.model.check_diffusers_installed", lambda: None) + + def fail_auto_config(*_args, **_kwargs): + raise AssertionError("local diffusion pipelines must be detected before AutoConfig") + + monkeypatch.setattr("transformers.AutoConfig.from_pretrained", fail_auto_config) + + assert is_diffusion_model(str(tmp_path)) is True + + +def test_missing_local_pipeline_component_does_not_probe_auto_config(monkeypatch, tmp_path): + component_path = Path(tmp_path, "feature_extractor") + + def fail_auto_config(*_args, **_kwargs): + raise AssertionError("missing local component paths must not be sent to AutoConfig") + + monkeypatch.setattr("transformers.AutoConfig.from_pretrained", fail_auto_config) + + assert is_diffusion_model(str(component_path)) is False From 21813e63e8c76a24003e3f38aaa2e07bb9e0d28d Mon Sep 17 00:00:00 2001 From: changwangss Date: Tue, 28 Jul 2026 18:15:33 +0800 Subject: [PATCH 8/9] docs: add SVDQuant details Signed-off-by: changwangss --- docs/svdquant_details.md | 415 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 415 insertions(+) create mode 100644 docs/svdquant_details.md diff --git a/docs/svdquant_details.md b/docs/svdquant_details.md new file mode 100644 index 0000000000..07707735e0 --- /dev/null +++ b/docs/svdquant_details.md @@ -0,0 +1,415 @@ +# SVDQuant — Details (Experimental) + +> **Experimental feature.** The currently validated end-to-end path targets +> MXFP4 quantization of Diffusers FLUX transformers. The core SVDQuant +> preprocessor is model-agnostic, but the `svdquant_nunchaku` format is +> runtime-loadable only for supported model adapters; FLUX is the first +> supported adapter. + +This document describes AutoRound's SVDQuant preprocessing, its interaction +with RTN and SignRound, activation-aware smoothing, and Nunchaku export. + +SVDQuant decomposes a linear weight into a low-rank branch and a quantized +residual branch: + +```text +W ~= Q(R) + U @ V + +Linear(x) ~= QuantizedLinear(x, Q(R)) + Linear(Linear(x, V), U) +``` + +The low-rank branch uses a configured BF16, FP16, or FP32 dtype during +quantization, while the larger residual branch is quantized to MXFP4. The +Nunchaku export materializes the low-rank tensors as BF16 by default. This +keeps a small, high-precision correction for weight directions that are +expensive to represent with four-bit values. + +--- + +## Prerequisites + +Use an isolated environment with a CUDA-enabled PyTorch build appropriate for +the target GPU. When working from the AutoRound source tree, install AutoRound +and the diffusion dependencies before quantization: + +```bash +pip install -e . +pip install -r test/test_cuda/requirements_diffusion.txt +``` + +Nunchaku is not needed to generate the quantized artifact. It is needed only +for inference and must be built or installed with MXFP4 support against a +compatible PyTorch/CUDA ABI. Verify the inference environment before loading +an exported pipeline: + +```bash +python -c "import torch, nunchaku; print(torch.__version__, torch.version.cuda, nunchaku.__file__)" +``` + +The local diffusion-caption dataset is a tab-separated file with `id` and +`caption` columns: + +```text +id\tcaption +0\tA photo of a cat +1\tA city street at night +``` + +--- + +## 1. Algorithm composition + +SVDQuant is a structural preprocessor, not a terminal quantizer. It must be +followed by exactly one block quantizer: + +```text +SVDQuantTransform -> RTN or SignRound -> output format +``` + +The CLI expresses this as an ordered algorithm list: + +| CLI | Pipeline | Calibration | +|-----|----------|-------------| +| `--algorithm svdquant,rtn` | SVDQuant + RTN | Data-free when smooth is disabled | +| `--algorithm svdquant,auto_round` | SVDQuant + SignRound | Uses SignRound calibration | + +The SVDQuant residual outer iteration always uses deployment-compatible RTN +QDQ. This is independent of the terminal quantizer: selecting SignRound means +SignRound optimizes the residual linears after SVDQuant has created them; it +does not replace RTN inside the residual decomposition loop. + +### 1.1 Processing flow + +```text +Original projection weights + | + +-- optional smooth search using bounded calibration calls + | + +-- group related FLUX projections + | (for example Q/K/V share one down factor) + | + +-- truncated SVD and residual outer iteration + | + +-- replace nn.Linear with SVDQuantLinear + | +-- MXFP4 residual_linear + | +-- BF16/FP16 lora_down + | +-- BF16/FP16 lora_up + | +-- optional smooth factor + | + +-- terminal RTN or SignRound quantization + | + +-- optional Nunchaku FLUX export +``` + +Related projections are decomposed as one stacked matrix where required by +the runtime format. For example, FLUX Q/K/V projections share one low-rank +down factor and retain separate up factors. + +--- + +## 2. Residual decomposition + +With `residual_iters=1`, one grouped weight matrix `W` is processed as follows: + +1. Compute a rank-`r` truncated SVD of `W`. +2. Materialize the low-rank factors in `low_rank_dtype`. +3. Form the residual `R = W - U @ V`. +4. Wrap `R`, `U`, and `V` in `SVDQuantLinear`. +5. Let the terminal RTN or SignRound stage quantize the residual linear. + +With `residual_iters > 1`, SVDQuant additionally runs a deployment-compatible +RTN QDQ outer loop before the terminal quantizer. Each iteration computes a +new low-rank decomposition from `W - Q(R_previous)`, forms and QDQs the new +residual, and evaluates `Q(R) + U @ V`. This alternates low-rank fitting and +residual quantization. `--enable-svdquant-residual-early-stop` stops after the +error first becomes worse than the best accepted candidate. + +The selection metric depends on smooth mode: + +- **Smooth disabled:** weight reconstruction squared error. +- **Smooth enabled:** calibration output squared error for the projection + group, including deployment-compatible residual and activation QDQ. + +Increasing `residual_iters` increases repeated SVD and QDQ work. It does not +change the terminal SignRound `--iters`; the two iteration counts control +different optimization loops. + +--- + +## 3. Activation-aware smooth search + +Smooth search is disabled by default. Enable it with: + +```bash +--enable-svdquant-smooth +``` + +For each FLUX projection group, AutoRound collects activation channel spans +and weight channel spans, then evaluates Alpha/Beta smooth factors. For +`G = --svdquant-smooth-num-grids`, the candidate set is: + +```text +(alpha, beta) = (0, 0) +(alpha, 0) for alpha in {1/G, ..., (G-1)/G} +(alpha, 1-alpha) for the same alpha values +``` + +The total candidate count is `1 + 2 * (G - 1)`. The default `G=20` evaluates +39 candidates per smooth group. Every candidate is scored against cached +floating-point outputs, and the lowest finite output error is selected. + +`--svdquant-smooth-max-calibration-calls` bounds the retained calls per smooth +group using reservoir sampling. Calls that are not retained are not copied to +the CPU cache. This bounds smooth-search memory, but the terminal SignRound +calibration flow still has its own sample and block-input requirements. + +Selected factor statistics are logged: + +```text +scale_min, scale_max, scale_ratio, below_1e-3, above_20 +``` + +Extreme-factor warnings are diagnostic. They do not modify or clamp the +selected factor. + +--- + +## 4. Configuration reference + +`SVDQuantConfig` is defined in +`auto_round/algorithms/transforms/svdquant/config.py`. + +| Python field | CLI option | Default | Description | +|--------------|------------|---------|-------------| +| `rank` | `--svdquant-rank` | `32` | Rank of the high-precision correction branch; Nunchaku export requires a positive rank. | +| `smooth_enabled` | `--enable-svdquant-smooth` | `False` | Enable activation-aware Alpha/Beta search. | +| `smooth_num_grids` | `--svdquant-smooth-num-grids` | `20` | Grid resolution; produces `1 + 2 * (G - 1)` candidates. | +| `smooth_max_calibration_calls` | `--svdquant-smooth-max-calibration-calls` | `128` | Maximum retained smooth calibration calls per group. | +| `smooth_eps` | Python API only | `1e-6` | Positive floor used while constructing factors. | +| `residual_iters` | `--svdquant-residual-iters` | `1` | Alternating low-rank/residual iterations. | +| `residual_early_stop` | `--enable-svdquant-residual-early-stop` | `False` | Stop when the selected error no longer improves. | +| `residual_quant_method` | `--svdquant-residual-quant-method` | `"rtn"` | Residual outer-loop QDQ method; currently fixed to RTN. | +| `low_rank_dtype` | `--svdquant-low-rank-dtype` | `"bf16"` | Low-rank factor dtype: BF16, FP16, or FP32 aliases. | +| `target_modules` | `--svdquant-target-modules` | `None` | Comma-separated module-name substrings to transform. | +| `exclude_modules` | `--svdquant-exclude-modules` | `None` | Comma-separated module-name substrings to keep untransformed. | +| `model_adapter` | `--svdquant-model-adapter` | `None` in Python; `"auto"` in CLI | Export mapping: `auto`, `flux`, or `identity`; `None` is resolved as auto. | + +All SVDQuant CLI options use hyphens. Existing shared AutoRound options retain +their existing spelling. + +--- + +## 5. CLI usage + +### 5.1 No-smooth SVDQuant + RTN + +This is the fastest data-free path: + +```bash +CUDA_VISIBLE_DEVICES=0 auto-round \ + --model /path/to/FLUX.1-dev \ + --model_dtype bf16 \ + --scheme MXFP4 \ + --algorithm svdquant,rtn \ + --disable_opt_rtn \ + --iters 0 \ + --nblocks 1 \ + --svdquant-rank 32 \ + --svdquant-residual-iters 20 \ + --enable-svdquant-residual-early-stop \ + --svdquant-residual-quant-method rtn \ + --svdquant-low-rank-dtype bf16 \ + --svdquant-model-adapter flux \ + --format svdquant_nunchaku \ + --device 0 \ + --low_gpu_mem_usage \ + --disable_low_cpu_mem_usage \ + --output_dir ./flux-dev-mxfp4-svdquant-rtn +``` + +### 5.2 Smooth SVDQuant + SignRound + +This path performs smooth search first, then runs SignRound on the residual +linears: + +```bash +CUDA_VISIBLE_DEVICES=0 auto-round \ + --model /path/to/FLUX.1-dev \ + --model_dtype bf16 \ + --scheme MXFP4 \ + --algorithm svdquant,auto_round \ + --iters 200 \ + --nblocks 1 \ + --nsamples 128 \ + --batch_size 1 \ + --num_inference_steps 50 \ + --dataset /path/to/coco2017-captions.tsv \ + --svdquant-rank 32 \ + --enable-svdquant-smooth \ + --svdquant-smooth-num-grids 20 \ + --svdquant-smooth-max-calibration-calls 128 \ + --svdquant-residual-iters 20 \ + --enable-svdquant-residual-early-stop \ + --svdquant-residual-quant-method rtn \ + --svdquant-low-rank-dtype bf16 \ + --svdquant-model-adapter flux \ + --format svdquant_nunchaku \ + --device 0 \ + --low_gpu_mem_usage \ + --disable_low_cpu_mem_usage \ + --output_dir ./flux-dev-mxfp4-svdquant-signround +``` + +This is a quality-oriented example, not a universal memory-safe preset. +Reduce `nsamples`, diffusion steps, smooth calls, or SignRound iterations for +workflow validation before running a full calibration. + +--- + +## 6. Python API + +Pass SVDQuant and one terminal quantizer through `alg_configs`: + +```python +from auto_round import AutoRound +from auto_round.algorithms.quantization.rtn.config import RTNConfig +from auto_round.algorithms.transforms.svdquant import SVDQuantConfig + +autoround = AutoRound( + "/path/to/FLUX.1-dev", + scheme="MXFP4", + model_dtype="bf16", + alg_configs=[ + SVDQuantConfig( + rank=32, + smooth_enabled=False, + residual_iters=20, + residual_early_stop=True, + model_adapter="flux", + ), + RTNConfig(disable_opt_rtn=True), + ], + device_map=0, + low_gpu_mem_usage=True, + low_cpu_mem_usage=False, +) + +autoround.quantize_and_save( + "./flux-dev-mxfp4-svdquant-rtn", + format="svdquant_nunchaku", +) +``` + +The example above covers the data-free RTN path. Use the CLI workflow in +section 5.2 for smooth SignRound calibration so the caption dataset, batch +size, and diffusion-step arguments are explicit. + +--- + +## 7. FLUX mapping and Nunchaku export + +`--format svdquant_nunchaku` currently requires: + +- `scheme=MXFP4` +- E2M1 weight and activation data +- group size 32 +- symmetric weight and activation quantization +- dynamic activation quantization +- a supported runtime model adapter (`flux` or auto-detected FLUX) + +The FLUX adapter transforms only runtime-supported projections. AdaNorm +linears are exported as W4A16 group-64 tensors, while RMSNorm weights and +top-level transformer tensors are exported as BF16. Non-transformer pipeline +components are saved through Diffusers in the dtype in which they were loaded; +the CLI examples load them with `--model_dtype bf16`. + +The output is a self-contained Diffusers pipeline: + +```text +output/ + model_index.json + scheduler/ + tokenizer/ + tokenizer_2/ + text_encoder/ + text_encoder_2/ + vae/ + transformer/ + config.json + diffusion_pytorch_model.safetensors +``` + +The transformer safetensors metadata contains `config`, +`quantization_config`, `model_class`, `comfy_config`, and `format`. +AutoRound's exporter does not import Nunchaku; Nunchaku is required only when +loading the exported model for inference. + +Export may take several minutes after quantization because fused FLUX records +are recomposed and decomposed to the configured common rank on CPU. + +--- + +## 8. Nunchaku inference + +After the prerequisite Nunchaku environment check succeeds, load the pipeline +directly: + +```python +import torch +from diffusers import FluxPipeline + +pipe = FluxPipeline.from_pretrained( + "./flux-dev-mxfp4-svdquant-signround", + torch_dtype=torch.bfloat16, + local_files_only=True, +) +pipe.enable_model_cpu_offload() + +generator = torch.Generator(device="cuda").manual_seed(12345) +image = pipe( + "A cat holding a sign that says Hello world", + num_inference_steps=20, + guidance_scale=3.5, + generator=generator, + height=512, + width=512, +).images[0] +image.save("flux-svdquant-mxfp4.png") +``` + +`model_index.json` points the transformer entry to +`nunchaku.NunchakuFluxTransformer2dModel`, so loading fails if a compatible +Nunchaku package is not installed. + +--- + +## 9. Memory and performance guidance + +- Start with a small end-to-end smoke configuration before a full calibration. +- `--low_gpu_mem_usage` enables component-level CPU offload for diffusion + calibration rather than placing the complete pipeline on one GPU. +- `--svdquant-smooth-max-calibration-calls` bounds only the retained smooth + evaluation pool; it does not replace `--nsamples` or SignRound calibration. +- More smooth candidates increase group replay time approximately linearly. +- More residual iterations repeat SVD and QDQ; early stop can reduce this work + only after the error becomes worse. +- More SignRound `--iters` increases terminal optimization time and is separate + from smooth and residual iteration costs. +- CPU memory, GPU memory, and runtime depend strongly on calibration samples, + diffusion steps, image size, and model placement. + +--- + +## 10. Current limitations + +- Runtime-loadable export currently supports FLUX through the FLUX adapter. +- The Nunchaku format is restricted to deployable MXFP4 E2M1 group-32 schemes. +- SVDQuant currently requires `nblocks=1`. +- Smooth calibration is output-aware and can be expensive because every + Alpha/Beta candidate replays retained group calls. +- Full-quality settings are hardware- and dataset-dependent; the examples are + starting points rather than guaranteed optimal presets. +- Export uses one common low-rank rank in Nunchaku metadata, so fused records + are recomputed at that rank instead of preserving a sum of source ranks. +- Generated-image quality must be evaluated over a representative prompt set; + a single smoke image validates loading and numerical stability only. From 1cd31043f6f0648093655ed6f2ec8e35613ce0d5 Mon Sep 17 00:00:00 2001 From: changwangss Date: Wed, 29 Jul 2026 09:09:48 +0800 Subject: [PATCH 9/9] fix: validate SVDQuant calibration config Signed-off-by: changwangss --- .../algorithms/transforms/svdquant/config.py | 12 ++++++++++ auto_round/autoround.py | 2 ++ .../algorithms/test_svdquant_residual.py | 6 +++++ test/test_cpu/core/test_pipeline_fail_fast.py | 24 +++++++++++++++++++ 4 files changed, 44 insertions(+) diff --git a/auto_round/algorithms/transforms/svdquant/config.py b/auto_round/algorithms/transforms/svdquant/config.py index 16719a6e40..931b9e8429 100644 --- a/auto_round/algorithms/transforms/svdquant/config.py +++ b/auto_round/algorithms/transforms/svdquant/config.py @@ -53,6 +53,18 @@ def __init__( ) if smooth_eps <= 0: raise ValueError(f"`smooth_eps` must be positive, got {smooth_eps!r}") + if not isinstance(low_rank_dtype, str) or low_rank_dtype.lower() not in { + "bf16", + "bfloat16", + "fp16", + "float16", + "fp32", + "float32", + }: + raise ValueError( + "`low_rank_dtype` must be one of bf16, bfloat16, fp16, float16, fp32, or float32, " + f"got {low_rank_dtype!r}" + ) if type(residual_iters) is not int or residual_iters < 1: raise ValueError(f"`residual_iters` must be a positive integer, got {residual_iters!r}") if type(residual_early_stop) is not bool: diff --git a/auto_round/autoround.py b/auto_round/autoround.py index 2b67300dec..3bb20e274e 100644 --- a/auto_round/autoround.py +++ b/auto_round/autoround.py @@ -244,8 +244,10 @@ def __new__( seed=seed, low_cpu_mem_usage=low_cpu_mem_usage, layer_config=layer_config, + dataset=dataset, nsamples=nsamples, seqlen=seqlen, + batch_size=batch_size, **entry_kwargs, ) diff --git a/test/test_cpu/algorithms/test_svdquant_residual.py b/test/test_cpu/algorithms/test_svdquant_residual.py index a1dc87dceb..c6ade9d94b 100644 --- a/test/test_cpu/algorithms/test_svdquant_residual.py +++ b/test/test_cpu/algorithms/test_svdquant_residual.py @@ -77,6 +77,7 @@ def test_svdquant_config_defaults_to_data_free_single_iteration(): ({"smooth_enabled": 1}, "smooth_enabled"), ({"smooth_num_grids": 1}, "smooth_num_grids"), ({"smooth_max_calibration_calls": 0}, "smooth_max_calibration_calls"), + ({"low_rank_dtype": "bf116"}, "low_rank_dtype"), ({"residual_iters": 0}, "residual_iters"), ({"residual_quant_method": "signround"}, "residual_quant_method"), ], @@ -86,6 +87,11 @@ def test_svdquant_config_rejects_invalid_structural_options(kwargs, field): SVDQuantConfig(**kwargs) +@pytest.mark.parametrize("dtype", ["bf16", "bfloat16", "fp16", "float16", "fp32", "float32"]) +def test_svdquant_config_accepts_supported_low_rank_dtype_aliases(dtype): + assert SVDQuantConfig(low_rank_dtype=dtype).low_rank_dtype == dtype + + def test_truncated_svd_returns_shared_down_factor_for_stacked_projection_group(): torch.manual_seed(0) qkv = torch.randn(12, 8, dtype=torch.float32) diff --git a/test/test_cpu/core/test_pipeline_fail_fast.py b/test/test_cpu/core/test_pipeline_fail_fast.py index f4e59bc1d6..6850c63a0c 100644 --- a/test/test_cpu/core/test_pipeline_fail_fast.py +++ b/test/test_cpu/core/test_pipeline_fail_fast.py @@ -142,6 +142,30 @@ def _fake_init(self, config, **kwargs): assert spinquant_cfg.trainable_smooth is rotation_config["trainable_smooth"] +def test_compat_entry_forwards_calibration_data_with_algorithm_configs(monkeypatch): + captured = {} + dataset = [{"text": "calibration prompt"}] + + def _fake_new_autoround(*args, **kwargs): + captured.update(kwargs) + return object() + + monkeypatch.setattr("auto_round.compressors.entry.AutoRound", _fake_new_autoround) + + from auto_round.autoround import AutoRound as CompatAutoRound + + CompatAutoRound( + "dummy-model", + scheme="W4A16", + alg_configs=[RTNConfig(disable_opt_rtn=True)], + dataset=dataset, + batch_size=1, + ) + + assert captured["dataset"] is dataset + assert captured["batch_size"] == 1 + + def test_entry_warns_and_drops_unsupported_kwargs(monkeypatch, tiny_opt_model_path): calls = []