Add experimental softmax reparameterization for output-head quantization

CHANGELOG.rst

@@ -13,6 +13,8 @@ Changelog
 
 *Quantization*
 
+- Add an experimental validation-KL search over softmax-equivalent linear output heads before weight-only quantization. See ``experimental/softmax_reparameterization`` for usage and current limitations.
+
 - Backfill checkpoint aliases for nine more published NVFP4 releases, so each is reachable from its source model's hub path: ``zai-org/GLM-5.1`` and ``GLM-5.2``, ``MiniMaxAI/MiniMax-M2.5`` and ``MiniMax-M3``, ``deepseek-ai/DeepSeek-V3.1`` and ``DeepSeek-V3.2``, ``Qwen/Qwen3-235B-A22B-Instruct-2507`` and ``-Thinking-2507``, and ``Qwen/Qwen3.6-27B``. Each imports an existing general or architecture recipe wholesale rather than copying its body.
 - Add composed Hugging Face AutoQuantize recipes that run fixed PTQ or weight AutoQuantize before
   a separate KV-cache AutoQuantize stage, with independent resumable checkpoints for the weight and

experimental/softmax_reparameterization/README.md

@@ -0,0 +1,100 @@
+# Softmax reparameterization for output-head quantization
+
+An experimental implementation of [Kadav et al., *Softmax Reparameterization for
+Output-Head Quantization*](https://arxiv.org/abs/2609.31291).
+
+For a vocabulary-by-hidden weight matrix `W`, each candidate is
+`W_t = W - t * W.mean(dim=0)`. Every vocabulary row receives the same vector
+subtraction, preserving the full-precision softmax distribution algebraically.
+Quantizing these equivalent heads produces different residuals. We fit the same
+ModelOpt quantizer independently for each candidate and select the coefficient
+with the lowest forward KL to the **original** head on validation states.
+
+## Usage
+
+Install ModelOpt from this checkout, then run the synthetic example:
+
+```bash
+python -m experimental.softmax_reparameterization.example --device cpu
+python -m experimental.softmax_reparameterization.example --device cuda --algorithm gptq
+```
+
+For an existing head, capture the final hidden states **entering the output
+projection**. Use articles disjoint across quantizer fitting, coefficient
+validation, and final test evaluation; exclude padding tokens. The caller must
+ensure these splits are disjoint. Each tensor has shape `[tokens, hidden_width]`
+and may reside on CPU to limit accelerator memory usage.
+
+```python
+import copy
+
+import torch
+
+import modelopt.torch.opt as mto
+import modelopt.torch.quantization as mtq
+from experimental.softmax_reparameterization import head_kl, search_head
+
+# source_head is an unquantized torch.nn.Linear immediately preceding softmax.
+config = copy.deepcopy(mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG)
+result = search_head(source_head, fit_states, validation_states, config)
+print(result.coefficient, result.validation_kl)
+print("Test KL:", head_kl(source_head, result.model, test_states))
+mto.save(result.model, "selected_head.pt")
+restored = mto.restore(torch.nn.Sequential(copy.deepcopy(source_head)), "selected_head.pt")
+```
+
+The returned model is a one-layer `Sequential`, with the head at index `0`; quantizer name patterns refer
+to this container, not to the original language model. The default blockwise INT4
+preset applies to it without changing ModelOpt's default LM-head exclusions.
+To use GPTQ, set `config["algorithm"] = "gptq"` before search. This is ModelOpt's
+calibration implementation; its rounding and fitting conventions need not match
+the paper's quantizers.
+
+The ordered default grid is `[-2, -1, -0.5, 0, 0.5, 1, 1.5, 2, 2.5, 3, 4, 5, 6, 8]`.
+Zero is required; one provides the fixed mean-centering comparison. Exact ties
+keep the first candidate. The full validation curve is returned, and the selected
+calibrated candidate is retained without refitting. Including zero guarantees
+non-increasing **measured validation KL**, not improvement on unseen data.
+
+## Scope and numerical behavior
+
+- Supports plain, unsharded `torch.nn.Linear` heads, with optional unchanged bias,
+  on CPU or a single CUDA device. Inputs and weights use FP32, FP16, or BF16;
+  log-softmax and KL are computed in FP32 and accumulated in FP64.
+- Supports weight-only fake quantization with `max` or `gptq` calibration.
+  Activation quantization, composite quantizers, nonlinear logit transformations
+  (including soft-capping), and grouped coefficient search are outside this PR.
+- Every candidate has independent weights and quantizer state. The source head
+  and any tied input embedding remain untouched, including on failure. Using the
+  selected head in a tied model requires a separate output weight allocation and
+  disabling weight tying when saving/reloading the full model.
+- Equivalence holds in exact arithmetic. Casting shifted weights and low-precision
+  matrix products can introduce small differences even before quantization.
+- The search runs the head only and does not retrain or rerun the decoder. Peak
+  memory includes the source, candidate, selected head, and temporary folded-weight
+  copies; GPTQ additionally needs a hidden-width-squared Hessian.
+
+## Model support and deployment
+
+This first integration tests synthetic linear heads. It does not claim that the
+paper's model-level results have been reproduced with ModelOpt. Capture and
+full-model adapters are intentionally left to the caller, who must verify a
+shift-compatible logit path. Replacing a Hugging Face child module also needs a
+model-level ModelOpt conversion/export workflow; the standalone head container
+must not be mistaken for an export-ready language-model checkpoint.
+
+The shift is stored in weights and requires no extra inference operator for a
+linear-softmax head. This implementation evaluates fake-quantized heads and tests
+ModelOpt head save/restore. Packed export and TensorRT-LLM/vLLM/SGLang serving
+are **not validated** by this prototype; no latency benefit is claimed here.
+
+## Tests
+
+```bash
+pytest tests/unit/torch/quantization/test_softmax_reparameterization.py -q
+pytest tests/gpu/torch/quantization/test_softmax_reparameterization_cuda.py -q
+```
+
+CPU tests compare selection with independent ModelOpt candidates, check softmax
+equivalence, source/tied-embedding isolation, tie handling, invalid inputs, and
+save/restore. CUDA tests exercise BF16 heads with INT4 max and GPTQ calibration.

experimental/softmax_reparameterization/__init__.py

@@ -0,0 +1,18 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""Experimental scalar search over equivalent linear-softmax output heads."""
+
+from .search import *

experimental/softmax_reparameterization/example.py

@@ -0,0 +1,58 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""Run a reproducible, synthetic head-only search without downloading a model."""
+
+import argparse
+import copy
+import json
+
+import torch
+from torch import nn
+
+import modelopt.torch.quantization as mtq
+from experimental.softmax_reparameterization import head_kl, search_head
+
+
+def main():
+    """Compare selected and unshifted heads on disjoint synthetic test states."""
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--device", choices=["cpu", "cuda"], default="cpu")
+    parser.add_argument("--algorithm", choices=["max", "gptq"], default="max")
+    args = parser.parse_args()
+    torch.manual_seed(0)
+    head = nn.Linear(128, 512, bias=False).to(args.device).eval()
+    with torch.no_grad():
+        head.weight.add_(torch.randn(128, device=args.device) * 0.1)
+    fit, validation, test = torch.randn(3, 64, 128, device=args.device).unbind()
+    config = copy.deepcopy(mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG)
+    config["algorithm"] = args.algorithm
+    result = search_head(head, fit, validation, config)
+    baseline = search_head(head, fit, validation, config, coefficients=(0,))
+    print(
+        json.dumps(
+            {
+                "coefficient": result.coefficient,
+                "validation_kl": result.validation_kl,
+                "test_kl_unshifted": head_kl(head, baseline.model, test),
+                "test_kl_selected": head_kl(head, result.model, test),
+            },
+            indent=2,
+        )
+    )
+
+
+if __name__ == "__main__":
+    main()

experimental/softmax_reparameterization/requirements.txt

@@ -0,0 +1 @@
+# No dependencies beyond base ModelOpt. Tests require pytest.

experimental/softmax_reparameterization/search.py

@@ -0,0 +1,157 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""Validation-selected output-head shifts (https://arxiv.org/abs/2609.31291)."""
+
+import copy
+import math
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from torch import nn
+
+import modelopt.torch.quantization as mtq
+from modelopt.torch.quantization.nn import TensorQuantizer
+
+__all__ = ["DEFAULT_COEFFICIENTS", "SearchResult", "head_kl", "search_head"]
+
+DEFAULT_COEFFICIENTS = (-2, -1, -0.5, 0, 0.5, 1, 1.5, 2, 2.5, 3, 4, 5, 6, 8)
+
+
+@dataclass
+class SearchResult:
+    """Selected ModelOpt head container and token-averaged validation KL for every candidate."""
+
+    model: nn.Sequential
+    coefficient: float
+    validation_kl: tuple[tuple[float, float], ...]
+
+
+def _check_states(states: torch.Tensor, width: int) -> None:
+    if states.ndim != 2 or states.shape[0] == 0 or states.shape[1] != width:
+        raise ValueError("States must be a nonempty [tokens, hidden_width] tensor.")
+    if not states.is_floating_point() or not torch.isfinite(states).all():
+        raise ValueError("States must contain finite floating-point values.")
+
+
+@torch.no_grad()
+def head_kl(
+    source: nn.Linear,
+    candidate: nn.Module,
+    states: torch.Tensor,
+    batch_size: int = 128,
+) -> float:
+    """Mean forward KL from source to candidate on final hidden states.
+
+    Both readouts use the source weight dtype; log-softmax and KL use FP32.
+    States may reside on CPU and are moved one batch at a time. Use disjoint
+    test states here after freezing the coefficient selected by ``search_head``.
+    """
+    _check_states(states, source.in_features)
+    if batch_size <= 0:
+        raise ValueError("batch_size must be positive.")
+    total = torch.zeros((), device=source.weight.device, dtype=torch.float64)
+    for batch in states.split(batch_size):
+        batch = batch.to(device=source.weight.device, dtype=source.weight.dtype)
+        log_p = source(batch).float().log_softmax(-1)
+        log_q = candidate(batch).float().log_softmax(-1)
+        total += (log_p.exp() * (log_p - log_q)).sum(-1).double().sum()
+    value = (total / states.shape[0]).item()
+    if not math.isfinite(value):
+        raise ValueError("Non-finite KL; check the head, quantizer, and hidden states.")
+    return value
+
+
+@torch.no_grad()
+def search_head(
+    head: nn.Linear,
+    fit_states: torch.Tensor,
+    validation_states: torch.Tensor,
+    quant_config: dict[str, Any],
+    coefficients: Sequence[float] = DEFAULT_COEFFICIENTS,
+    batch_size: int = 128,
+) -> SearchResult:
+    """Fit shifted copies of a linear-softmax head and select by validation KL.
+
+    Each candidate subtracts ``t * head.weight.float().mean(0)`` from every
+    vocabulary row before independently fitting the same ModelOpt quantizer.
+    The original head (including a shared embedding parameter) is never mutated.
+
+    Args:
+        head: Plain, unquantized ``nn.Linear`` immediately preceding softmax.
+            Nonlinear logit paths and distributed/sharded heads are not supported.
+        fit_states: Final hidden states used exclusively for quantizer fitting.
+        validation_states: States from disjoint articles used only for selection.
+            The caller is responsible for enforcing article-level disjointness.
+        quant_config: ModelOpt weight-only fake-quant config for ``max`` or ``gptq``.
+            Patterns address a one-layer Sequential (e.g. ``0.weight_quantizer``).
+        coefficients: Finite, unique ordered candidates including zero; ties keep
+            the first candidate. Including one provides a fixed-centering control.
+        batch_size: Maximum number of token states per readout.
+
+    Returns:
+        The selected calibrated Sequential, coefficient, and validation curve.
+        Keep the container when using ModelOpt save/restore. Selection guarantees
+        no worse measured validation KL than t=0, not improvement on unseen data.
+    """
+    if type(head) is not nn.Linear:
+        raise TypeError("head must be a plain, unquantized nn.Linear.")
+    if head.weight.device.type not in ("cpu", "cuda"):
+        raise ValueError("Only materialized CPU or CUDA heads are supported.")
+    if head.weight.dtype not in (torch.float32, torch.float16, torch.bfloat16):
+        raise ValueError("Use FP32, FP16, or BF16 head weights.")
+    values = tuple(float(t) for t in coefficients)
+    if not values or not all(math.isfinite(t) for t in values):
+        raise ValueError("coefficients must be nonempty and finite.")
+    if 0.0 not in values or len(set(values)) != len(values):
+        raise ValueError("coefficients must be unique and include zero.")
+    if batch_size <= 0:
+        raise ValueError("batch_size must be positive.")
+    _check_states(fit_states, head.in_features)
+    _check_states(validation_states, head.in_features)
+    algorithm = quant_config.get("algorithm", "max")
+    method = algorithm.get("method") if isinstance(algorithm, dict) else algorithm
+    if method not in ("max", "gptq"):
+        raise ValueError("This prototype supports max and gptq calibration only.")
+
+    def fit(model):
+        for batch in fit_states.split(batch_size):
+            model(batch.to(device=head.weight.device, dtype=head.weight.dtype))
+
+    mean = head.weight.float().mean(0)
+    best_model, best_t, best_kl = None, 0.0, math.inf
+    curve = []
+    for coefficient in values:
+        candidate = nn.Sequential(copy.deepcopy(head)).eval().requires_grad_(False)
+        if coefficient != 0:
+            candidate[0].weight.copy_(head.weight.float() - coefficient * mean)
+        mtq.quantize(candidate, copy.deepcopy(quant_config), forward_loop=fit)
+        quantizer = candidate[0].weight_quantizer
+        if not isinstance(quantizer, TensorQuantizer) or not quantizer.is_enabled:
+            raise ValueError("Enable a single weight quantizer for the head.")
+        if not quantizer.fake_quant:
+            raise ValueError("Search requires fake quantization; pack only after selection.")
+        if candidate[0].input_quantizer.is_enabled or candidate[0].output_quantizer.is_enabled:
+            raise ValueError("Search supports weight-only quantization.")
+        # Avoid requantizing the full vocabulary matrix for each validation batch.
+        with mtq.temporarily_fold_weights(candidate):
+            kl = head_kl(head, candidate, validation_states, batch_size)
+        curve.append((coefficient, kl))
+        if kl < best_kl:
+            best_model, best_t, best_kl = candidate, coefficient, kl
+    assert best_model is not None
+    return SearchResult(best_model, best_t, tuple(curve))

tests/gpu/torch/quantization/test_softmax_reparameterization_cuda.py

@@ -0,0 +1,43 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""CUDA validation with ModelOpt INT4 max and GPTQ calibration."""
+
+import copy
+
+import pytest
+import torch
+from torch import nn
+
+import modelopt.torch.quantization as mtq
+from experimental.softmax_reparameterization import head_kl, search_head
+
+
+@pytest.mark.parametrize("algorithm", ["max", "gptq"])
+def test_cuda_head_search(algorithm):
+    torch.manual_seed(7)
+    head = nn.Linear(128, 256, bias=False, device="cuda", dtype=torch.bfloat16).eval()
+    original = head.weight.detach().clone()
+    fit, val, test = torch.randn(3, 64, 128, device="cuda", dtype=torch.bfloat16).unbind()
+    config = copy.deepcopy(mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG)
+    config["algorithm"] = algorithm
+    result = search_head(head, fit, val, config, (0, 1, 4), batch_size=32)
+    curve = dict(result.validation_kl)
+    assert curve[result.coefficient] <= curve[0]
+    assert head_kl(head, result.model, val, 64) == pytest.approx(
+        curve[result.coefficient], abs=2e-5
+    )
+    assert head_kl(head, result.model, test) >= -1e-6
+    torch.testing.assert_close(head.weight, original, rtol=0, atol=0)