diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 7960dac388b..3bdb6e70698 100755 --- a/CHANGELOG.rst +++ b/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 `_) 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/``. diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index 344292264fa..4618a0b3e8a 100644 --- a/modelopt/torch/quantization/config.py +++ b/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 diff --git a/modelopt/torch/quantization/nn/functional.py b/modelopt/torch/quantization/nn/functional.py index 94cbcf74fda..2a5f78f1817 100644 --- a/modelopt/torch/quantization/nn/functional.py +++ b/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 diff --git a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py index c50804dd2a9..60a7b9ee6f5 100644 --- a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py +++ b/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) diff --git a/tests/gpu/torch/quantization/test_hadamard.py b/tests/gpu/torch/quantization/test_hadamard.py index 1173eba2097..7da6a3ca122 100644 --- a/tests/gpu/torch/quantization/test_hadamard.py +++ b/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( diff --git a/tests/unit/torch/quantization/test_tensor_quant_cpu.py b/tests/unit/torch/quantization/test_tensor_quant_cpu.py index 218acd76f5d..b21375ef8e4 100644 --- a/tests/unit/torch/quantization/test_tensor_quant_cpu.py +++ b/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},