Add TensorQuantizer Random Hadamard rotation seed

CHANGELOG.rst

@@ -17,6 +17,8 @@ Changelog
 
 **New Features**
 
+- 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.
+- Add ``rotate.seed`` to torch quantizer configs. The default ``None`` keeps regular Hadamard rotation; an integer seed enables deterministic Random Hadamard Transform signs.
 - Add the **D-PACE** loss objective for DFlash speculative-decoding training (`arXiv:2605.18810 <https://arxiv.org/abs/2605.18810>`_) and make it the default (``dflash_loss_objective: dpace``). It replaces the static exponential position decay with dynamic, confidence-derived per-position weights that adapt to whichever block positions currently limit acceptance. Smoothing is controlled by ``dflash_dpace_alpha`` (default 0.5); set ``dflash_loss_objective: decay`` to restore the previous static schedule. Training-only and detached from the gradient (no architecture or inference change).
 - Add the ``day0-release`` agent skill (``.agents/skills/day0-release/``), a deterministic end-to-end driver that chains the PTQ → evaluation → comparison skills (the evaluation stage deploys the checkpoint itself) with an enforced gate after each stage and returns a publish decision (ACCEPT / REGRESSION / ANOMALOUS / INFEASIBLE). Ships three GPU-free, unit-tested gate scripts (``gate_ptq.py``, ``gate_run.py``, ``gate_compare.py``) that validate checkpoint coverage, evaluation-run completeness, and baseline-vs-candidate accuracy threshold. v1 reports and stops on regression; the recipe-search loop is deferred.
 - Add **streaming** speculative-decoding training (EAGLE3 / DFlash): the draft trains on base-model hidden states produced on the fly by a co-located ``vllm serve`` (no disk dump), moved trainer-side over NIXL RDMA, scaling to multi-node (dedicated serve replicas + DDP trainers). New launcher examples for NVFP4 Kimi-K2.5 / K2.6 on GB200/aarch64 under ``tools/launcher/examples/moonshotai/``.

modelopt/torch/quantization/config.py

@@ -287,8 +287,10 @@ class RotateConfig(ModeloptBaseConfig):
     """
 
     enable: bool = False
+    mode: Literal["rotate", "rotate_back"] = "rotate"
     rotate_fp32: bool = False
     block_size: int | None = None
+    seed: int | None = None
 
     @field_validator("block_size", mode="before")
     @classmethod
@@ -298,6 +300,14 @@ def validate_block_size(cls, v):
             raise ValueError(f"block_size must be a positive int, got {v!r}")
         return v
 
+    @field_validator("seed", mode="before")
+    @classmethod
+    def validate_seed(cls, v):
+        """Validate seed is a non-negative int (mode=before to catch bool before int coercion)."""
+        if v is not None and (isinstance(v, bool) or not isinstance(v, int) or v < 0):
+            raise ValueError(f"seed must be a non-negative int, got {v!r}")
+        return v
+
 
 class QuantizerAttributeConfig(ModeloptBaseConfig):
     """Quantizer attribute type."""
@@ -328,7 +338,7 @@ class QuantizerAttributeConfig(ModeloptBaseConfig):
     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
@@ -337,14 +347,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/nn/functional.py

@@ -16,6 +16,7 @@
 """Some supportive functions."""
 
 import warnings
+from functools import lru_cache
 
 import torch
 from torch.autograd import Function
@@ -98,7 +99,16 @@ def _largest_pow2_divisor(n: int) -> int:
     return n & (-n)
 
 
-def normalized_hadamard_transform(inputs, rotate_fp32=False, block_size=None):
+@lru_cache(maxsize=16)
+def _random_signs(seed: int, dim: int, device: torch.device, dtype: torch.dtype):
+    generator = torch.Generator(device=device).manual_seed(seed)
+    signs = torch.randint(0, 2, (dim,), generator=generator, device=device, dtype=torch.int8)
+    return signs.to(dtype=dtype).mul_(2).sub_(1)
+
+
+def normalized_hadamard_transform(
+    inputs, rotate_fp32=False, block_size=None, random_sign_seed=None, inverse=False
+):
     """Normalized fast hadamard transform.
 
     Supports block-granular RHT for dimensions that are not a power of 2.
@@ -112,6 +122,9 @@ def normalized_hadamard_transform(inputs, rotate_fp32=False, block_size=None):
         block_size: Block size for block-granular RHT. Must be power of 2 and divide
             inputs.shape[-1]. If None: use full-dimension FHT when dim is power of 2;
             otherwise auto-select the largest power-of-2 divisor of the dimension.
+        random_sign_seed: If set, apply a deterministic random sign diagonal before
+            Hadamard. The same seed is used for the inverse transform.
+        inverse: If True, apply inverse RHT order for seeded transforms.
 
     Returns:
         Rotated tensor with same shape as inputs.
@@ -129,6 +142,13 @@ def normalized_hadamard_transform(inputs, rotate_fp32=False, block_size=None):
     dtype = inputs.dtype
     if rotate_fp32:
         inputs = inputs.to(torch.float32)
+    signs = (
+        _random_signs(random_sign_seed, dim, inputs.device, inputs.dtype)
+        if random_sign_seed is not None
+        else None
+    )
+    if signs is not None and not inverse:
+        inputs = inputs * signs
 
     if block_size is None and utils.is_pow2(dim):
         # Full-dimension FHT (original behavior)
@@ -161,4 +181,7 @@ def normalized_hadamard_transform(inputs, rotate_fp32=False, block_size=None):
         )
         outputs = rotated.reshape(inputs.shape)
 
+    if signs is not None and inverse:
+        outputs = outputs * signs
+
     return outputs.to(dtype) if rotate_fp32 else outputs

modelopt/torch/quantization/nn/modules/tensor_quantizer.py

@@ -619,6 +619,33 @@ def rotate_block_size(self):
             return self._rotate.get("block_size", None)
         return None
 
+    @property
+    def rotate_seed(self):
+        """Seed for deterministic RHT signs, or None for normalized HT."""
+        if isinstance(self._rotate, RotateConfig):
+            return self._rotate.seed if self._rotate.enable else None
+        if isinstance(self._rotate, dict) and self.rotate_is_enabled:
+            return self._rotate.get("seed", 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 _rotate_inputs(self, inputs, inverse=False):
+        return normalized_hadamard_transform(
+            inputs,
+            rotate_fp32=self.rotate_is_fp32,
+            block_size=self.rotate_block_size,
+            random_sign_seed=self.rotate_seed,
+            inverse=inverse,
+        )
+
     def disable_calib(self):
         """Disable calibration."""
         self._if_calib = False
@@ -1090,13 +1117,12 @@ 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
@@ -1159,6 +1185,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, inverse=True)
+
         return outputs
 
     def _short_amax(self, fmt=".2e"):
@@ -1209,9 +1238,12 @@ def extra_repr(self):
             else ""
         )
         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})"
+        if self.rotate_seed is not None:
+            s += f" (seed={self.rotate_seed})"
         s += (
             f" calibrator={self._calibrator.__class__.__name__}"
             if (self._calibrator is not None)

tests/gpu/torch/quantization/test_hadamard.py

@@ -48,6 +48,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 +69,22 @@ 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("block_size", [None, 4])
+def test_hadamard_transform_seeded_inverse(block_size):
+    x = torch.rand(4, 8, device="cuda")
+    x_h = normalized_hadamard_transform(x, block_size=block_size, random_sign_seed=123)
+    x_roundtrip = normalized_hadamard_transform(
+        x_h, block_size=block_size, random_sign_seed=123, inverse=True
+    )
+
+    assert torch.allclose(x_roundtrip, x, rtol=1e-5, atol=1e-6)
+    assert torch.allclose(
+        x_h, normalized_hadamard_transform(x, block_size=block_size, random_sign_seed=123)
+    )
 
 
 @pytest.mark.parametrize(

tests/unit/torch/quantization/test_tensor_quant_cpu.py

@@ -22,8 +22,14 @@
 from _test_utils.torch.quantization.tensor_quant_common import FakeTensorQuantTester
 
 import modelopt.torch.quantization as mtq
+import modelopt.torch.quantization.nn.modules.tensor_quantizer as tensor_quantizer_module
 from modelopt.torch.quantization.config import QuantizerAttributeConfig
-from modelopt.torch.quantization.nn import TensorQuantizer
+from modelopt.torch.quantization.nn import (
+    TensorQuantizer,
+    register_quant_backend,
+    unregister_quant_backend,
+)
+from modelopt.torch.quantization.nn.functional import _random_signs
 
 
 class TestFakeTensorQuantCPU(FakeTensorQuantTester):
@@ -56,6 +62,28 @@ 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,
+                "seed": 123,
+            }
+        )
+
+        assert quant_attr_cfg.model_dump(exclude_unset=True)["rotate"] == {
+            "enable": True,
+            "mode": "rotate_back",
+            "rotate_fp32": True,
+            "block_size": 8,
+            "seed": 123,
+        }
+
+        with pytest.raises(ValueError, match="seed must be a non-negative int"):
+            QuantizerAttributeConfig(rotate={"enable": True, "seed": -1})
+
     def test_num_bits(self):
         """Test num_bits for both integer and tuple cases."""
 
@@ -92,6 +120,96 @@ 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, random_sign_seed=None, inverse=False):
+        calls.append((rotate_fp32, block_size, random_sign_seed, inverse))
+        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")
+
+
+def test_tensor_quantizer_rotate_mode_preserves_default_path(monkeypatch):
+    outputs, inputs, calls, quantizer = _run_rotated_backend(monkeypatch, rotate={"enable": True})
+
+    assert not quantizer.rotate_back_is_enabled
+    assert torch.equal(outputs, (inputs + 10) * 2)
+    assert calls == [(False, None, None, False)]
+
+
+def test_tensor_quantizer_rotate_mode_can_rotate_back(monkeypatch):
+    outputs, inputs, calls, quantizer = _run_rotated_backend(
+        monkeypatch,
+        rotate={
+            "enable": True,
+            "mode": "rotate_back",
+            "rotate_fp32": True,
+            "block_size": 8,
+            "seed": 123,
+        },
+    )
+
+    assert quantizer.rotate_back_is_enabled
+    assert torch.equal(outputs, ((inputs + 10) * 2) + 10)
+    assert calls == [(True, 8, 123, False), (True, 8, 123, True)]
+
+
+def test_tensor_quantizer_rotate_back_rejects_real_quant(monkeypatch):
+    def fail_if_rotated(
+        inputs, rotate_fp32=False, block_size=None, random_sign_seed=None, inverse=False
+    ):
+        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]]))
+
+
+def test_random_signs_are_cached_by_seed_dim_device_and_dtype():
+    _random_signs.cache_clear()
+
+    try:
+        signs = _random_signs(123, 8, torch.device("cpu"), torch.float32)
+        cached_signs = _random_signs(123, 8, torch.device("cpu"), torch.float32)
+        fp16_signs = _random_signs(123, 8, torch.device("cpu"), torch.float16)
+        different_seed_signs = _random_signs(124, 8, torch.device("cpu"), torch.float32)
+        different_dim_signs = _random_signs(123, 16, torch.device("cpu"), torch.float32)
+
+        assert cached_signs.data_ptr() == signs.data_ptr()
+        assert torch.equal(cached_signs, signs)
+        assert fp16_signs.dtype == torch.float16
+        assert fp16_signs.data_ptr() != signs.data_ptr()
+        assert different_seed_signs.data_ptr() != signs.data_ptr()
+        assert different_dim_signs.data_ptr() != signs.data_ptr()
+    finally:
+        _random_signs.cache_clear()
+
+
 WINT4INT8_CFG = {
     "quant_cfg": [
         {"quantizer_name": "*", "enable": False},