Add TensorQuantizer rotate-back mode

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/<model>/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**
 

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):

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

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)

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

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

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)