diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 32d5a3388e2..48aed5eb8b9 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -42,6 +42,7 @@ Changelog - Add **Domino** speculative-decoding training: the parallel DFlash draft backbone plus a lightweight GRU causal correction head, selected via ``dflash_architecture_config.projector_type=domino``. Trained with a base/final dual loss whose ``dflash_lambda_base_start``/``dflash_lambda_base_decay_ratio`` curriculum decays the base-loss weight 1→0. Exports in the z-lab drafter format; recipe at ``modelopt_recipes/general/speculative_decoding/domino.yaml``. Training only — the inference path is not wired up yet. - Add Torch-TensorRT FP8 deployment example for HuggingFace ViT (``examples/torch_trt/``): ``torch_tensorrt_ptq.py`` covers ``mtq.quantize`` → ``torch_tensorrt.compile(ir="dynamo")``, and ``torch_tensorrt_accuracy.py`` reports the compiled model's ImageNet-1k top-1/top-5 accuracy via the ``onnx_ptq`` ``evaluate`` harness (the unquantized baseline is Torch-TensorRT-compiled too, for an apples-to-apples comparison). Ships a ViT-tuned FP8 PTQ recipe under ``modelopt_recipes/huggingface/vit/ptq/`` (``fp8.yaml``) composed from the shared ``modelopt_recipes/configs/`` units: it quantizes the encoder Linears, patch-embed ``nn.Conv2d``, ``classifier``, and per-block LayerNorm inputs plus the attention Q/K/V BMMs and softmax. Verified on ``google/vit-base-patch16-224`` (ImageNet-1k 50k validation): FP8 stays within 0.13 pp Top-1 of the FP16 baseline. - Add **AutoQuantize recipe** support: ``mtq.auto_quantize`` can be driven declaratively from a YAML recipe (``RecipeType.AUTO_QUANTIZE`` / ``AutoQuantizeConfig``) specifying candidate formats, the ``effective_bits`` target, cost model (incl. ``active_moe`` and ``excluded_module_name_patterns``), scoring method, and disabled layers. Adds an ``effective_bits`` cost-model override on ``QuantizeConfig`` / ``QuantizerAttributeConfig`` (block-scale-accurate NVFP4 = 4.5 via ``configs/numerics/nvfp4``). Shipped recipes live under ``modelopt_recipes/general/auto_quantize/`` and model-specific ones under ``modelopt_recipes/huggingface//auto_quantize/``. +- Add ``rotate.mode`` to torch quantizer configs. The default ``"rotate"`` keeps the existing rotate-before-quantize behavior; ``"rotate_back"`` enables fake-quant rotate → quantize → rotate-back for TensorQuantizer. **Bug Fixes** diff --git a/modelopt/torch/export/plugins/vllm_fakequant_hf.py b/modelopt/torch/export/plugins/vllm_fakequant_hf.py index 8883964daac..acb1968e070 100644 --- a/modelopt/torch/export/plugins/vllm_fakequant_hf.py +++ b/modelopt/torch/export/plugins/vllm_fakequant_hf.py @@ -27,7 +27,6 @@ import torch.nn as nn import modelopt.torch.opt as mto -from modelopt.torch.quantization.config import RotateConfig from modelopt.torch.quantization.conversion import quantizer_state from modelopt.torch.quantization.model_calib import enable_stats_collection, finish_stats_collection from modelopt.torch.quantization.nn import QuantModule, SequentialQuantizer, TensorQuantizer @@ -137,15 +136,6 @@ def _check_all_weight_quantizers_disabled(model: nn.Module) -> None: ) -def disable_rotate(quantizer: TensorQuantizer): - """Return a disabled copy of the quantizer's ``_rotate`` field, preserving its type.""" - if isinstance(quantizer._rotate, RotateConfig): - return RotateConfig(enable=False) - if isinstance(quantizer._rotate, dict): # backward compat: old checkpoints stored a dict - return dict(quantizer._rotate, enable=False) - return False - - def _fakequant_fused_experts_weights( module: nn.Module, module_name: str, @@ -638,16 +628,12 @@ def export_hf_vllm_fq_checkpoint( if isinstance(quantizer, SequentialQuantizer): quantizer.disable() for sub in quantizer: - orig_rotate = sub._rotate - if sub.rotate_is_enabled: - sub._rotate = disable_rotate(sub) - wqs_to_restore.append((sub, orig_rotate)) + wqs_to_restore.append((sub, sub._rotate)) + sub.disable_rotate() elif isinstance(quantizer, TensorQuantizer): quantizer.disable() - orig_rotate = quantizer._rotate - if quantizer.rotate_is_enabled: - quantizer._rotate = disable_rotate(quantizer) - wqs_to_restore.append((quantizer, orig_rotate)) + wqs_to_restore.append((quantizer, quantizer._rotate)) + quantizer.disable_rotate() quantizer_state_dict = get_quantizer_state_dict(model) for key in list(quantizer_state_dict): diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index f4a935239a8..44ae162c474 100644 --- a/modelopt/torch/quantization/config.py +++ b/modelopt/torch/quantization/config.py @@ -286,9 +286,29 @@ class RotateConfig(ModeloptBaseConfig): for transform details. """ - enable: bool = False - rotate_fp32: bool = False - block_size: int | None = None + enable: bool = ModeloptField( + default=False, + title="Enable input rotation.", + description="If True, applies a normalized Hadamard transform before quantization.", + ) + mode: Literal["rotate", "rotate_back"] = ModeloptField( + default="rotate", + title="Rotation mode.", + description=( + "Use 'rotate' for input rotation only, or 'rotate_back' to apply the transform " + "again after fake quantization." + ), + ) + rotate_fp32: bool = ModeloptField( + default=False, + title="Run rotation in float32.", + description="If True, computes the rotation in float32 before casting back to the input dtype.", + ) + block_size: int | None = ModeloptField( + default=None, + title="Rotation block size.", + description="Positive block size for block-wise rotation, or None to rotate the full input.", + ) @field_validator("block_size", mode="before") @classmethod @@ -344,7 +364,7 @@ def _validate_effective_bits(cls, v: float | None) -> float | None: def validate_config(cls, values): """Validate quantizer config.""" - def _validate_recursive(value): + def _validate_recursive(value, field_name=None): """Recursively validate config structure.""" if value is None: return @@ -353,14 +373,16 @@ def _validate_recursive(value): for item in value: _validate_recursive(item) elif isinstance(value, dict): + if field_name == "rotate": + return if len(value) == 1 and "enable" in value and value["enable"] is True: raise ValueError( "Invalid quantizer config: Cannot specify only {'enable': True}. " "Additional parameters are required when enabling quantization." ) # Recurse into nested dicts - for v in value.values(): - _validate_recursive(v) + for k, v in value.items(): + _validate_recursive(v, k) _validate_recursive(values) return values diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index 7dbdd36d04e..1adda0511a6 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -619,7 +619,11 @@ def print_quant_summary(model: nn.Module, output_dir: str | None = None): def fold_weight(model: nn.Module, keep_attrs: bool = False): - """Fold weight quantizer for fast evaluation.""" + """Fold weight quantizer for fast evaluation. + + Any weight-quantizer rotation is folded into the weights and disabled so subsequent + forwards do not re-rotate the already-folded weights. + """ for name, module in model.named_modules(): if isinstance(module, QuantModule): module.fold_weight(keep_attrs) diff --git a/modelopt/torch/quantization/nn/functional.py b/modelopt/torch/quantization/nn/functional.py index 94cbcf74fda..533e7ebf68d 100644 --- a/modelopt/torch/quantization/nn/functional.py +++ b/modelopt/torch/quantization/nn/functional.py @@ -73,26 +73,6 @@ def backward(ctx, grad_output): clip = ClipFunction.apply -class FastHadamardTransform(Function): - """The fast Hadamard transform. - - This only works for inputs.shape[-1] == power of 2. - """ - - @staticmethod - def forward(ctx, inputs): - """Hadamard forward.""" - assert utils.is_pow2(inputs.shape[-1]), ( - "Fast hadamard only works for inputs.shape[-1] == power of 2." - ) - return fast_hadamard_transform.hadamard_transform(inputs) # type: ignore[name-defined] - - @staticmethod - def backward(ctx, grad_outputs): - """Hadamard backward.""" - return fast_hadamard_transform.hadamard_transform(grad_outputs) # type: ignore[name-defined] - - def _largest_pow2_divisor(n: int) -> int: """Return the largest power of 2 that divides n.""" return n & (-n) @@ -130,35 +110,29 @@ def normalized_hadamard_transform(inputs, rotate_fp32=False, block_size=None): if rotate_fp32: inputs = inputs.to(torch.float32) - if block_size is None and utils.is_pow2(dim): - # Full-dimension FHT (original behavior) - outputs = FastHadamardTransform.apply(inputs) / torch.sqrt( - torch.tensor(dim, dtype=torch.float32) - ) - else: - # Block-granular RHT - if block_size is None: + # Full-dimension FHT is just block-granular RHT with block_size == dim. + if block_size is None: + if utils.is_pow2(dim): + block_size = dim + else: block_size = _largest_pow2_divisor(dim) if block_size < 2: raise RuntimeError( f"Block RHT: dimension {dim} has no power-of-2 divisor >= 2. " "Set rotate.block_size explicitly (e.g. 128) or use a dimension divisible by a power of 2." ) - if not utils.is_pow2(block_size): - raise ValueError(f"Block RHT: block_size must be power of 2, got {block_size}.") - if dim % block_size != 0: - raise RuntimeError( - f"Block RHT: inputs.shape[-1]={dim} is not divisible by block_size={block_size}. " - f"Use a block_size that divides {dim} (e.g. {_largest_pow2_divisor(dim)})." - ) - n_blocks = dim // block_size - # Reshape to (..., n_blocks, block_size) - flat = inputs.reshape(-1, dim) - blocks = flat.reshape(-1, n_blocks, block_size) - # Apply FHT per block (last dim) - rotated = FastHadamardTransform.apply(blocks) / torch.sqrt( - torch.tensor(block_size, dtype=torch.float32) + if not utils.is_pow2(block_size): + raise ValueError(f"Block RHT: block_size must be power of 2, got {block_size}.") + if dim % block_size != 0: + raise RuntimeError( + f"Block RHT: inputs.shape[-1]={dim} is not divisible by block_size={block_size}. " + f"Use a block_size that divides {dim} (e.g. {_largest_pow2_divisor(dim)})." ) - outputs = rotated.reshape(inputs.shape) + + # hadamard_transform is autograd-aware and fuses the normalization scale into the kernel. + rotated = fast_hadamard_transform.hadamard_transform( # type: ignore[name-defined] + inputs.reshape(-1, dim // block_size, block_size).contiguous(), scale=block_size**-0.5 + ) + outputs = rotated.reshape(inputs.shape) return outputs.to(dtype) if rotate_fp32 else outputs diff --git a/modelopt/torch/quantization/nn/modules/quant_embedding.py b/modelopt/torch/quantization/nn/modules/quant_embedding.py index 5004c11b3c1..690afd8a334 100644 --- a/modelopt/torch/quantization/nn/modules/quant_embedding.py +++ b/modelopt/torch/quantization/nn/modules/quant_embedding.py @@ -40,6 +40,9 @@ class _QuantEmbedding(QuantModule): table (weight) and the lookup output (an activation feeding downstream layers) are quantizable. + TODO: Remove the example-side ``*embed_tokens*quantizer`` rotation workaround + once rotation configs stop targeting embedding token/input paths. + Quantizer roles: - ``weight_quantizer``: quantizes the embedding table (``self.weight``). - ``input_quantizer``: permanently disabled placeholder diff --git a/modelopt/torch/quantization/nn/modules/quant_module.py b/modelopt/torch/quantization/nn/modules/quant_module.py index 419c6f4924f..9c9aee478a8 100644 --- a/modelopt/torch/quantization/nn/modules/quant_module.py +++ b/modelopt/torch/quantization/nn/modules/quant_module.py @@ -17,6 +17,7 @@ import contextlib import warnings +from collections.abc import Iterable from typing import Any import torch @@ -127,8 +128,36 @@ def iter_weights_for_calibration(self): weight_quantizer = getattr(self, quantizer_attr_names(weight_name).weight_quantizer) yield getattr(self, weight_name), weight_quantizer + @staticmethod + @torch.no_grad() + def _fold_weight_quantizer( + quantizer: TensorQuantizer, + weights: Iterable[torch.Tensor], + keep_attrs: bool = False, + ): + """Fold ``quantizer`` into each weight view in place, then disable and clean it once. + + ``weights`` is an iterable so a single quantizer shared across views (e.g. one + per-tensor quantizer over all experts of a fused MoE weight) can be folded view by + view while disabling and dropping its calibration attrs exactly once. + """ + for weight in weights: + weight.data.copy_(quantizer(weight.float().contiguous()).to(weight.dtype)) + quantizer.disable() + quantizer.disable_rotate() + if not keep_attrs: + for attr_name in ("_pre_quant_scale", "_amax"): + if hasattr(quantizer, attr_name): + delattr(quantizer, attr_name) + def fold_weight(self, keep_attrs: bool = False): - """Fold the weight for faster eval.""" + """Bake each fake-quant weight quantizer into its weight for faster eval. + + Every fake-quant weight quantizer is folded regardless of its enabled state. The folded + transform is baked into the stored weight and then disabled, so subsequent forwards use + the stored weight directly. Calibration buffers (``_pre_quant_scale``, ``_amax``) are + dropped unless ``keep_attrs``. + """ # Handle all attributes that end with _weight_quantizer for name in dir(self): attr = getattr(self, name) @@ -144,16 +173,7 @@ def fold_weight(self, keep_attrs: bool = False): f"{name} doesn't have a corresponding {weight_name} in {self.__class__.__name__}" ) weight = getattr(self, weight_name) - weight.data.copy_(attr(weight.float()).to(weight.dtype)) - attr.disable() - if not keep_attrs: - _attrs = [ - "_pre_quant_scale", - "_amax", - ] - for attr_name in _attrs: - if hasattr(attr, attr_name): - delattr(attr, attr_name) + self._fold_weight_quantizer(attr, (weight,), keep_attrs) QuantModuleRegistry = _DMRegistryCls("Quant", QuantModule) diff --git a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py index c50804dd2a9..c7649f9383d 100644 --- a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py +++ b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py @@ -266,6 +266,9 @@ def _block_sizes_setter(val): ) setattr(self, _tq_attribute_name, _setter(val)) + if isinstance(attribute_cfg, dict) and attribute_cfg == {"enable": False}: + self.disable_rotate() + if self.is_mx_format: self._pass_through_bwd = True @@ -619,6 +622,35 @@ def rotate_block_size(self): return self._rotate.get("block_size", None) return None + @property + def rotate_back_is_enabled(self): + """Check if inverse rotation should be applied after quantization.""" + if isinstance(self._rotate, RotateConfig): + return self._rotate.enable and self._rotate.mode == "rotate_back" + if isinstance(self._rotate, dict) and self.rotate_is_enabled: + return self._rotate.get("mode", "rotate") == "rotate_back" + return False + + def disable_rotate(self): + """Disable rotation while preserving the ``_rotate`` field's type. + + Idempotent. Used after folding a weight quantizer so the baked-in rotation is not + re-applied on subsequent forwards. + """ + if isinstance(self._rotate, RotateConfig): + self._rotate = self._rotate.model_copy(update={"enable": False}) + elif isinstance(self._rotate, dict): # backward compat: old checkpoints stored a dict + self._rotate = dict(self._rotate, enable=False) + else: + self._rotate = False + + def _rotate_inputs(self, inputs): + return normalized_hadamard_transform( + inputs, + rotate_fp32=self.rotate_is_fp32, + block_size=self.rotate_block_size, + ) + def disable_calib(self): """Disable calibration.""" self._if_calib = False @@ -1090,19 +1122,22 @@ def forward(self, inputs): if self.pre_quant_scale is not None: inputs = inputs * self.pre_quant_scale + if self.rotate_back_is_enabled and self._if_quant and not self.fake_quant: + raise ValueError("rotate_back mode is only supported with fake_quant=True.") + # Rotating the input if self.rotate_is_enabled: - inputs = normalized_hadamard_transform( - inputs, - rotate_fp32=self.rotate_is_fp32, - block_size=self.rotate_block_size, - ) + inputs = self._rotate_inputs(inputs) if self._disabled: # if quantizer is disabled, we still need to track the input dtype for saving the model # TODO: This is a temporary solution and needs to be removed once megatron supports # non-homogeneous layers self._input_dtype = inputs.dtype if hasattr(inputs, "dtype") else None + # Even when quantization is disabled, honor rotate_back so a rotate/rotate_back + # pair stays a no-op roundtrip instead of leaving the tensor rotated. + if self.rotate_back_is_enabled: + inputs = self._rotate_inputs(inputs) return inputs if ( @@ -1159,6 +1194,9 @@ def forward(self, inputs): if self.is_static_block_quant: outputs = self._reset_to_original_shape(outputs) + if self.rotate_back_is_enabled and isinstance(outputs, torch.Tensor): + outputs = self._rotate_inputs(outputs) + return outputs def _short_amax(self, fmt=".2e"): @@ -1185,6 +1223,14 @@ def _short_tensor(self, tensor: torch.Tensor, fmt=".2e"): return f"{tensor.item():{fmt}}" return f"[{tensor.min().item():{fmt}}, {tensor.max().item():{fmt}}]({tensor.numel()})" + def _rotation_extra_repr(self): + s = " rotated" if self.rotate_is_enabled else "" + s += " (rotate_back)" if self.rotate_back_is_enabled else "" + s += " (fp32)" if self.rotate_is_fp32 else "" + if self.rotate_block_size is not None: + s += f" (block={self.rotate_block_size})" + return s + def extra_repr(self): """Set the extra information about this module.""" if self._disabled: @@ -1194,7 +1240,8 @@ def extra_repr(self): if self.pre_quant_scale is not None else "" ) - return "disabled" + s += self._rotation_extra_repr() + return s s = f"{'unsigned ' if self._unsigned else ''}{self._num_bits} bit" s += " narrow" if (self._narrow_range) else "" s += " fake" if (self._fake_quant) else "" @@ -1208,10 +1255,7 @@ def extra_repr(self): if self.pre_quant_scale is not None else "" ) - s += " rotated" if self.rotate_is_enabled else "" - s += " (fp32)" if self.rotate_is_fp32 else "" - if self.rotate_block_size is not None: - s += f" (block={self.rotate_block_size})" + s += self._rotation_extra_repr() s += ( f" calibrator={self._calibrator.__class__.__name__}" if (self._calibrator is not None) @@ -1511,6 +1555,7 @@ class SequentialQuantizer(nn.Sequential): _delegated_methods = [ "reset_amax", "disable", + "disable_rotate", "enable", "load_calib_amax", "load_calib_bias", diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 6779c3f9ade..c86c90eaa59 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -1091,31 +1091,17 @@ def iter_weights_for_calibration(self): yield weight[idx], q def fold_weight(self, keep_attrs: bool = False): - """Fold per-expert weight quantizers into the fused 3-D weights. + """Bake each per-expert weight quantizer into its slice of the fused 3-D weight. - The base ``fold_weight`` only handles singular ``*_weight_quantizer`` - attributes. Fused experts use ``nn.ModuleList`` of per-expert quantizers - (``_weight_quantizers``, ``down_proj_weight_quantizers``), - which would otherwise be skipped, leaving ``_amax`` on every quantizer. + The base ``fold_weight`` only handles singular ``*_weight_quantizer`` attributes and + would skip the ``nn.ModuleList`` of per-expert quantizers used here. The per-expert + ``(weight_slice, quantizer)`` pairs are the same ones :meth:`iter_weights_for_calibration` + yields, so we reuse it; each fake-quant quantizer's quantization and rotation are folded + in and disabled, and calibration buffers are dropped unless ``keep_attrs``. """ - for weight_name, quantizers_name in ( - (self._first_proj_attr, self._first_proj_weight_quantizers_attr), - ("down_proj", "down_proj_weight_quantizers"), - ): - weight = getattr(self, weight_name, None) - quantizers = getattr(self, quantizers_name, None) - if weight is None or quantizers is None: - continue - for idx, q in enumerate(quantizers): - if not (isinstance(q, TensorQuantizer) and q.fake_quant): - continue - slice_ = weight.data[idx] - slice_.copy_(q(slice_.float()).to(weight.dtype)) - q.disable() - if not keep_attrs: - for attr_name in ("_pre_quant_scale", "_amax"): - if hasattr(q, attr_name): - delattr(q, attr_name) + for weight_slice, q in self.iter_weights_for_calibration(): + if isinstance(q, TensorQuantizer) and q.fake_quant: + self._fold_weight_quantizer(q, (weight_slice,), keep_attrs) class _QuantNonGatedFusedExperts(_QuantFusedExperts): diff --git a/modelopt/torch/quantization/plugins/vllm.py b/modelopt/torch/quantization/plugins/vllm.py index 95ca3240b73..aa6141dd48b 100644 --- a/modelopt/torch/quantization/plugins/vllm.py +++ b/modelopt/torch/quantization/plugins/vllm.py @@ -462,20 +462,13 @@ def forward(self, hidden_states: torch.Tensor, router_logits: torch.Tensor): @torch.no_grad() def fold_weight(self, keep_attrs: bool = False): # the MoE weights can be super large, it consumes too much memory, so we need to fold the weight one by one - for i in range(self.w13_weight.shape[0]): - self.w13_weight[i].copy_( - self.w13_weight_quantizer(self.w13_weight[i].float().contiguous()).to( - self.w13_weight.dtype - ) - ) - self.w13_weight_quantizer.disable() - for i in range(self.w2_weight.shape[0]): - self.w2_weight[i].copy_( - self.w2_weight_quantizer(self.w2_weight[i].float().contiguous()).to( - self.w2_weight.dtype - ) + for weight, quantizer in ( + (self.w13_weight, self.w13_weight_quantizer), + (self.w2_weight, self.w2_weight_quantizer), + ): + self._fold_weight_quantizer( + quantizer, (weight[i] for i in range(weight.shape[0])), keep_attrs ) - self.w2_weight_quantizer.disable() if torch.cuda.is_available(): torch.cuda.empty_cache() diff --git a/tests/gpu/torch/quantization/test_hadamard.py b/tests/gpu/torch/quantization/test_hadamard.py index 1173eba2097..d426a3e2ffa 100644 --- a/tests/gpu/torch/quantization/test_hadamard.py +++ b/tests/gpu/torch/quantization/test_hadamard.py @@ -30,11 +30,13 @@ from _test_utils.torch.quantization.models import SDPAAttention import modelopt.torch.quantization as mtq +from modelopt.torch.quantization.config import QuantizerAttributeConfig from modelopt.torch.quantization.conversion import ( set_quantizer_by_cfg, set_quantizer_by_cfg_context, ) from modelopt.torch.quantization.nn.functional import normalized_hadamard_transform +from modelopt.torch.quantization.nn.modules.tensor_quantizer import TensorQuantizer @pytest.mark.parametrize( @@ -48,6 +50,9 @@ def test_hadamard_transform(dim): xxt_h = x_h @ x_h.T # The numerical error can be large, especially for 16-bit floats. assert torch.allclose(xxt_h, xxt, atol=0.05) + x_roundtrip = normalized_hadamard_transform(x_h) + assert torch.allclose(x_roundtrip, x, rtol=1e-5, atol=1e-6) + x_h_fp32 = normalized_hadamard_transform(x, rotate_fp32=True) xxt_h_fp32 = x_h_fp32 @ x_h_fp32.T assert torch.allclose(xxt_h_fp32, xxt, atol=0.05) @@ -66,6 +71,8 @@ def test_hadamard_transform_block(dim, block_size): # Use rtol instead of atol: float32 accumulated error scales with value magnitude, # which grows with dim. 1e-3 relative tolerance is appropriate for float32 block RHT. assert torch.allclose(xxt_h, xxt, rtol=1e-3, atol=1e-6) + x_roundtrip = normalized_hadamard_transform(x_h, block_size=block_size) + assert torch.allclose(x_roundtrip, x, rtol=1e-5, atol=1e-6) @pytest.mark.parametrize( @@ -104,3 +111,24 @@ def test_kv_rotate(rotate_fp32): assert not torch.allclose(output_ref, output_test1, atol=0.05) mtq.unregister(SDPAAttention) + + +@pytest.mark.parametrize( + "mode", + ["rotate", "rotate_back"], +) +def test_rotate_backward(mode): + """Autograd backward should flow through a rotation-enabled TensorQuantizer.""" + x = torch.randn(4, 64, device="cuda", requires_grad=True) + quantizer = TensorQuantizer( + QuantizerAttributeConfig(num_bits=8, axis=None, rotate={"enable": True, "mode": mode}), + amax=x.detach().abs().amax(), + ).cuda() + + out = quantizer(x) + out.sum().backward() + + assert x.grad is not None + assert x.grad.shape == x.shape + assert torch.isfinite(x.grad).all() + assert not torch.all(x.grad == 0) diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index 1829777e87f..9fa836bb620 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -25,8 +25,10 @@ from _test_utils.torch.quantization.tied_modules import tie_fused_experts_3d_params import modelopt.torch.quantization as mtq +import modelopt.torch.quantization.nn.modules.tensor_quantizer as tensor_quantizer_module from modelopt.torch.export.moe_utils import _export_fused_experts from modelopt.torch.export.quant_utils import get_quant_config, get_quantization_format +from modelopt.torch.quantization.config import QuantizerAttributeConfig from modelopt.torch.quantization.conversion import _normalize_fused_experts_quantizer_name from modelopt.torch.quantization.model_calib import local_hessian_calibrate from modelopt.torch.quantization.nn import QuantModuleRegistry, TensorQuantizer @@ -336,6 +338,49 @@ def test_expert_index_recovery(self): assert recovered_idx == idx, f"Expected {idx}, got {recovered_idx}" self._cleanup_registry(expert_type) + def _make_rotated_fused_experts(self, monkeypatch): + monkeypatch.setattr( + tensor_quantizer_module, + "normalized_hadamard_transform", + lambda inputs, rotate_fp32=False, block_size=None: inputs, + ) + model = _TinyMoEModel() + expert_type = type(model.moe.experts) + self._cleanup_registry(expert_type) + register_fused_experts_on_the_fly(model) + converted = QuantModuleRegistry.convert(model.moe.experts) + expert_quantizers = list(converted.gate_up_proj_weight_quantizers) + list( + converted.down_proj_weight_quantizers + ) + for q in expert_quantizers: + q.set_from_attribute_config( + QuantizerAttributeConfig(num_bits=8, rotate={"enable": True}) + ) + q.amax = torch.tensor(6.0) + return converted, expert_quantizers, expert_type + + def test_fold_weight_disables_per_expert_quantizers_and_rotation(self, monkeypatch): + converted, expert_quantizers, expert_type = self._make_rotated_fused_experts(monkeypatch) + try: + converted.fold_weight() + for q in expert_quantizers: + assert not q.is_enabled + assert not q.rotate_is_enabled + assert not hasattr(q, "_amax") + finally: + self._cleanup_registry(expert_type) + + def test_fold_weight_keep_attrs_keeps_amax_disables_rotation(self, monkeypatch): + converted, expert_quantizers, expert_type = self._make_rotated_fused_experts(monkeypatch) + try: + converted.fold_weight(keep_attrs=True) + for q in expert_quantizers: + assert not q.is_enabled + assert not q.rotate_is_enabled + assert hasattr(q, "_amax") + finally: + self._cleanup_registry(expert_type) + # --------------------------------------------------------------------------- # Tests for export diff --git a/tests/unit/torch/quantization/test_print.py b/tests/unit/torch/quantization/test_print.py index 56a351e799e..d6e3e0d3181 100644 --- a/tests/unit/torch/quantization/test_print.py +++ b/tests/unit/torch/quantization/test_print.py @@ -20,6 +20,7 @@ from modelopt.torch.quantization import calib, tensor_quant from modelopt.torch.quantization import nn as qnn +from modelopt.torch.quantization.config import QuantizerAttributeConfig from modelopt.torch.quantization.nn.modules.tensor_quantizer import TensorQuantizer @@ -32,6 +33,25 @@ def test_print_tensor_quantizer(self): test_quantizer = TensorQuantizer() print(test_quantizer) + def test_disabled_tensor_quantizer_repr_shows_enabled_state(self): + test_quantizer = TensorQuantizer( + QuantizerAttributeConfig( + enable=False, + rotate={ + "enable": True, + "mode": "rotate_back", + "rotate_fp32": True, + "block_size": 8, + }, + ) + ) + test_quantizer.pre_quant_scale = torch.tensor([1.0, 2.0]) + + assert test_quantizer.extra_repr() == ( + "disabled pre_quant_scale=[1.00e+00, 2.00e+00](2)" + " rotated (rotate_back) (fp32) (block=8)" + ) + def test_print_module(self): class _TestModule(nn.Module): def __init__(self): diff --git a/tests/unit/torch/quantization/test_quant_embedding.py b/tests/unit/torch/quantization/test_quant_embedding.py index 5a3d2535dcf..d7053d1e09f 100644 --- a/tests/unit/torch/quantization/test_quant_embedding.py +++ b/tests/unit/torch/quantization/test_quant_embedding.py @@ -108,6 +108,42 @@ def test_wildcard_config_keeps_input_quantizer_disabled(self): # Forward still works — input_quantizer is disabled and never applied. qemb(torch.randint(0, VOCAB_SIZE, (4, 6))) + def test_disable_update_clears_hard_disabled_input_quantizer_rotate_state(self): + qemb = _make_quant_embedding() + set_quantizer_attributes_partial( + qemb, + "*input_quantizer", + {"enable": True, "num_bits": 4, "rotate": {"enable": True, "mode": "rotate_back"}}, + ) + assert not qemb.input_quantizer.is_enabled + assert qemb.input_quantizer.num_bits == 4 + assert qemb.input_quantizer.rotate_is_enabled + assert qemb.input_quantizer.extra_repr() == "disabled rotated (rotate_back)" + + set_quantizer_attributes_partial(qemb, "*input_quantizer", {"enable": False}) + + assert not qemb.input_quantizer.is_enabled + assert qemb.input_quantizer.num_bits == 4 + assert not qemb.input_quantizer.rotate_is_enabled + assert qemb.input_quantizer.extra_repr() == "disabled" + + def test_disable_update_clears_weight_quantizer_rotate_state(self): + qemb = _make_quant_embedding() + set_quantizer_attributes_partial( + qemb, + "*weight_quantizer", + {"enable": True, "num_bits": 4, "rotate": {"enable": True, "mode": "rotate_back"}}, + ) + assert qemb.weight_quantizer.is_enabled + assert qemb.weight_quantizer.rotate_is_enabled + assert "rotated (rotate_back)" in qemb.weight_quantizer.extra_repr() + + set_quantizer_attributes_partial(qemb, "*weight_quantizer", {"enable": False}) + + assert not qemb.weight_quantizer.is_enabled + assert not qemb.weight_quantizer.rotate_is_enabled + assert "rotated" not in qemb.weight_quantizer.extra_repr() + # Export-path tests for QuantEmbedding live in tests/gpu/torch/export/test_export_embedding.py # because _export_quantized_weight bottoms out in torch.cuda.empty_cache(), which raises on diff --git a/tests/unit/torch/quantization/test_tensor_quant_cpu.py b/tests/unit/torch/quantization/test_tensor_quant_cpu.py index 218acd76f5d..ba352ec2162 100644 --- a/tests/unit/torch/quantization/test_tensor_quant_cpu.py +++ b/tests/unit/torch/quantization/test_tensor_quant_cpu.py @@ -22,8 +22,15 @@ from _test_utils.torch.quantization.tensor_quant_common import FakeTensorQuantTester import modelopt.torch.quantization as mtq -from modelopt.torch.quantization.config import QuantizerAttributeConfig -from modelopt.torch.quantization.nn import TensorQuantizer +import modelopt.torch.quantization.nn.modules.tensor_quantizer as tensor_quantizer_module +from modelopt.torch.quantization import QuantModuleRegistry +from modelopt.torch.quantization.config import QuantizerAttributeConfig, RotateConfig +from modelopt.torch.quantization.nn import ( + SequentialQuantizer, + TensorQuantizer, + register_quant_backend, + unregister_quant_backend, +) class TestFakeTensorQuantCPU(FakeTensorQuantTester): @@ -56,6 +63,18 @@ def test_from_to_dict(self, verbose): quant_attr_cfg_2 = QuantizerAttributeConfig(**quant_attr_cfg_1.dict()) assert quant_attr_cfg_1 == quant_attr_cfg_2 + def test_rotate_mode_serialization(self): + quant_attr_cfg = QuantizerAttributeConfig( + rotate={"enable": True, "mode": "rotate_back", "rotate_fp32": True, "block_size": 8} + ) + + assert quant_attr_cfg.model_dump(exclude_unset=True)["rotate"] == { + "enable": True, + "mode": "rotate_back", + "rotate_fp32": True, + "block_size": 8, + } + def test_num_bits(self): """Test num_bits for both integer and tuple cases.""" @@ -92,6 +111,227 @@ def test_num_bits(self): QuantizerAttributeConfig(enable=True, num_bits=(-1, 2)) +def _run_rotated_backend(monkeypatch, rotate): + calls = [] + + def rotate_fn(inputs, rotate_fp32=False, block_size=None): + calls.append((rotate_fp32, block_size)) + return inputs + 10 + + def backend(inputs, _tq): + return inputs * 2 + + monkeypatch.setattr(tensor_quantizer_module, "normalized_hadamard_transform", rotate_fn) + register_quant_backend("test_rotate_mode_backend", backend) + try: + quantizer = TensorQuantizer( + QuantizerAttributeConfig(rotate=rotate, backend="test_rotate_mode_backend") + ) + inputs = torch.tensor([[1.0, 2.0]]) + return quantizer(inputs), inputs, calls, quantizer + finally: + unregister_quant_backend("test_rotate_mode_backend") + + +@pytest.mark.parametrize( + ("rotate", "rotate_back_enabled", "expected_calls", "expected_fn"), + [ + ({"enable": True}, False, [(False, None)], lambda inputs: (inputs + 10) * 2), + ( + {"enable": True, "mode": "rotate_back", "rotate_fp32": True, "block_size": 8}, + True, + [(True, 8), (True, 8)], + lambda inputs: ((inputs + 10) * 2) + 10, + ), + ], +) +def test_tensor_quantizer_rotate_modes( + monkeypatch, rotate, rotate_back_enabled, expected_calls, expected_fn +): + outputs, inputs, calls, quantizer = _run_rotated_backend( + monkeypatch, + rotate=rotate, + ) + + assert quantizer.rotate_back_is_enabled is rotate_back_enabled + assert torch.equal(outputs, expected_fn(inputs)) + assert calls == expected_calls + + +def test_tensor_quantizer_rotate_back_rejects_real_quant(monkeypatch): + def fail_if_rotated(inputs, rotate_fp32=False, block_size=None): + raise AssertionError("rotate_back with fake_quant=False should fail before rotation") + + monkeypatch.setattr( + tensor_quantizer_module, + "normalized_hadamard_transform", + fail_if_rotated, + ) + quantizer = TensorQuantizer( + QuantizerAttributeConfig( + num_bits=8, + fake_quant=False, + rotate={"enable": True, "mode": "rotate_back"}, + ) + ) + + with pytest.raises(ValueError, match="rotate_back mode is only supported with fake_quant=True"): + quantizer(torch.tensor([[1.0, 2.0]])) + + +@pytest.mark.parametrize( + ("rotate", "rotate_back_enabled", "expected_call_count", "expected_fn"), + [ + ({"enable": True}, False, 1, lambda inputs: inputs + 10), + ({"enable": True, "mode": "rotate_back"}, True, 2, lambda inputs: inputs + 20), + ], +) +def test_tensor_quantizer_disabled_rotate_modes_roundtrip( + monkeypatch, rotate, rotate_back_enabled, expected_call_count, expected_fn +): + calls = [] + + def rotate_fn(inputs, rotate_fp32=False, block_size=None): + calls.append((rotate_fp32, block_size)) + return inputs + 10 + + monkeypatch.setattr(tensor_quantizer_module, "normalized_hadamard_transform", rotate_fn) + quantizer = TensorQuantizer(QuantizerAttributeConfig(rotate=rotate, enable=False)) + inputs = torch.tensor([[1.0, 2.0]]) + + outputs = quantizer(inputs) + + assert quantizer.rotate_back_is_enabled is rotate_back_enabled + assert torch.equal(outputs, expected_fn(inputs)) + assert len(calls) == expected_call_count + + +def test_disable_only_update_clears_regular_quantizer_rotate_state(): + quantizer = TensorQuantizer( + QuantizerAttributeConfig(rotate={"enable": True, "mode": "rotate_back", "block_size": 8}) + ) + assert quantizer.rotate_is_enabled + assert quantizer.rotate_back_is_enabled + + quantizer.set_from_attribute_config({"enable": False}) + + assert not quantizer.is_enabled + assert not quantizer.rotate_is_enabled + assert isinstance(quantizer._rotate, RotateConfig) + assert quantizer._rotate.mode == "rotate_back" + assert quantizer._rotate.block_size == 8 + + +def test_disable_rotate_preserves_type(): + # RotateConfig: enable off, other fields retained. + quantizer = TensorQuantizer( + QuantizerAttributeConfig(rotate={"enable": True, "mode": "rotate_back", "block_size": 8}) + ) + assert isinstance(quantizer._rotate, RotateConfig) + quantizer.disable_rotate() + assert isinstance(quantizer._rotate, RotateConfig) + assert quantizer._rotate.enable is False + assert quantizer._rotate.mode == "rotate_back" + assert quantizer._rotate.block_size == 8 + assert not quantizer.rotate_is_enabled + quantizer.disable_rotate() # idempotent + assert quantizer._rotate.enable is False + + # Raw dict (old checkpoints). + quantizer._rotate = {"enable": True, "mode": "rotate", "block_size": 4} + quantizer.disable_rotate() + assert quantizer._rotate == {"enable": False, "mode": "rotate", "block_size": 4} + + # Bool. + quantizer._rotate = True + quantizer.disable_rotate() + assert quantizer._rotate is False + + +def test_sequential_quantizer_disable_rotate_delegates(): + q0 = TensorQuantizer(QuantizerAttributeConfig(rotate={"enable": True})) + q1 = TensorQuantizer(QuantizerAttributeConfig(rotate={"enable": True, "mode": "rotate_back"})) + seq = SequentialQuantizer(q0, q1) + + seq.disable_rotate() + + assert not q0.rotate_is_enabled + assert not q1.rotate_is_enabled + + +def _make_qlinear_with_backend(monkeypatch, calls, backend_name, rotate=False): + def rotate_fn(inputs, rotate_fp32=False, block_size=None): + calls.append((rotate_fp32, block_size)) + return inputs + 10 + + def backend(inputs, _tq): + return inputs * 2 + + monkeypatch.setattr(tensor_quantizer_module, "normalized_hadamard_transform", rotate_fn) + register_quant_backend(backend_name, backend) + qlinear = QuantModuleRegistry.convert(torch.nn.Linear(4, 3)) + qlinear.input_quantizer.disable() + qlinear.output_quantizer.disable() + qlinear.weight_quantizer.set_from_attribute_config( + QuantizerAttributeConfig(rotate=rotate, backend=backend_name) + ) + return qlinear + + +@pytest.mark.parametrize( + ("rotate", "expected_weight_fn"), + [ + ({"enable": True}, lambda weight: (weight + 10) * 2), + ({"enable": True, "mode": "rotate_back"}, lambda weight: ((weight + 10) * 2) + 10), + ], +) +def test_fold_weight_disables_quantizer_without_extra_transform( + monkeypatch, rotate, expected_weight_fn +): + calls = [] + backend_name = "test_fold_backend" + qlinear = _make_qlinear_with_backend(monkeypatch, calls, backend_name, rotate=rotate) + try: + qlinear.weight_quantizer.amax = torch.tensor(1.0) + weight0 = qlinear.weight.detach().clone() + x = torch.randn(2, 4) + out_before = qlinear(x) + + qlinear.fold_weight() + + assert torch.allclose(qlinear.weight, expected_weight_fn(weight0)) + assert not qlinear.weight_quantizer.is_enabled + assert not qlinear.weight_quantizer.rotate_is_enabled + assert not hasattr(qlinear.weight_quantizer, "_amax") + + calls_after_fold = len(calls) + out_after = qlinear(x) + assert torch.allclose(out_after, out_before) + assert len(calls) == calls_after_fold + + # Second fold is a no-op: quantizer is disabled. + weight_after = qlinear.weight.detach().clone() + qlinear.fold_weight() + assert torch.allclose(qlinear.weight, weight_after) + finally: + unregister_quant_backend(backend_name) + + +def test_fold_weight_keep_attrs_keeps_amax(monkeypatch): + calls = [] + backend_name = "test_fold_backend_keep" + qlinear = _make_qlinear_with_backend(monkeypatch, calls, backend_name) + try: + qlinear.weight_quantizer.amax = torch.tensor(1.0) + + qlinear.fold_weight(keep_attrs=True) + + assert hasattr(qlinear.weight_quantizer, "_amax") + assert not qlinear.weight_quantizer.is_enabled + finally: + unregister_quant_backend(backend_name) + + WINT4INT8_CFG = { "quant_cfg": [ {"quantizer_name": "*", "enable": False},