Add experimental softmax reparameterization for output-head quantization

[asimkadav]

What does this PR do?

Type of change: new feature (experimental).

Low-bit quantization can assign different prediction errors to functionally equivalent output heads. This adds a head-only implementation of Softmax Reparameterization for Output-Head Quantization: subtract a scalar multiple of the vocabulary-row mean, independently fit the same ModelOpt quantizer for each candidate, and select by forward KL to the original head on disjoint validation states.

The implementation supports plain linear-softmax heads and weight-only max/gptq calibration. It preserves the original head and tied embeddings, returns the full validation curve, and retains the selected calibrated artifact without refitting. Zero is required in the candidate grid, guaranteeing no worse measured validation KL than the unshifted candidate; held-out improvements are not guaranteed.

This is an opt-in experimental addition. It changes no default quantization configs or LM-head exclusions. Nonlinear logit paths, grouped search, full-model adapters, packed export, and serving benchmarks are outside this PR. The paper's reported results should not be interpreted as ModelOpt measurements.

Usage

import copy
import modelopt.torch.quantization as mtq
from experimental.softmax_reparameterization import head_kl, search_head

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(head_kl(source_head, result.model, test_states))

States are final hidden vectors shaped [tokens, hidden_width]; the caller enforces disjoint articles across fitting, validation, and test. The returned one-layer Sequential supports ModelOpt save/restore. A runnable synthetic example is included.

Testing

• CPU: python -m pytest -o addopts='' tests/unit/torch/quantization/test_softmax_reparameterization.py -q — 3 passed.
• Tests use real ModelOpt quantization: independent candidate agreement, exact-arithmetic softmax equivalence within numerical tolerance, source/tied-embedding isolation, first-candidate tie handling, input validation, and ModelOpt save/restore.
• All pre-commit hooks passed on changed files.
• CPU synthetic example: held-out KL 0.00803 unshifted vs 0.000885 selected. These are synthetic smoke-test values, not language-model benchmarks.
• NVIDIA H100 80GB: combined CPU and CUDA BF16 max/GPTQ tests — 5 passed in 56.35 seconds (python -m pytest -o addopts='' tests/unit/torch/quantization/test_softmax_reparameterization.py tests/gpu/torch/quantization/test_softmax_reparameterization_cuda.py -q).
• Both CUDA examples completed over the full 14-candidate grid. Synthetic held-out KL: max 0.011655 → 0.000848; GPTQ 0.033649 → 0.002440. Both selected t=1. These validate execution and selection on synthetic heads, not the paper's language-model results.

Before your PR is "Ready for review"

• Is this change backward compatible?: ✅ Opt-in experimental module.
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md?: N/A. New implementation; no additional runtime dependencies.
• Did you write any new necessary tests?: ✅
• Did you update Changelog?: ✅
• Did you get Claude approval on this PR?: N/A; not yet requested.

Additional Information

Paper: https://arxiv.org/abs/2609.31291

CPU and H100 validation completed. Maintainer feedback on the experimental API is welcome. Packed deployment support is not claimed. NVIDIA-hosted CI requires maintainer authorization for this fork PR.

Summary by CodeRabbit

• New Features

  • Added an experimental search that selects an output-head reparameterization based on validation KL before weight-only quantization.
  • Added a runnable synthetic example to compare the selected result with an unshifted baseline.

• Documentation

  • Documented usage, supported configurations, input requirements, limitations, and test commands.

• Tests

  • Added CPU and CUDA coverage for candidate selection, validation results, and preservation of the original model.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

Navigate logical layers of code changes, visualize relationships, and explore their blast radius.

🧰 Additional context used

📚 Code guidelines (1)

.github/PULL_REQUEST_TEMPLATE.md — configured

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: Repository: NVIDIA/Model-Optimizer/.coderabbit.yaml
• Review profile: CHILL
• Plan: Enterprise
• Run ID: aaedfadc-082e-4448-b6ee-fec0e46fa8a2

📥 Commits

Reviewing files that changed from the base of the PR and between 6a10687 and e8bf828.

📒 Files selected for processing (8)

• CHANGELOG.rst
• experimental/softmax_reparameterization/README.md
• experimental/softmax_reparameterization/__init__.py
• experimental/softmax_reparameterization/example.py
• experimental/softmax_reparameterization/requirements.txt
• experimental/softmax_reparameterization/search.py
• tests/gpu/torch/quantization/test_softmax_reparameterization_cuda.py
• tests/unit/torch/quantization/test_softmax_reparameterization.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.

---

📝 Walkthrough

Walkthrough

Adds an experimental search for softmax-equivalent linear output-head shifts. The search independently quantizes each candidate, measures validation KL against the original head, and returns the candidate with the lowest KL. Documentation, a synthetic example, and CPU and CUDA tests accompany the implementation.

Changes

Output-Head Search

Layer / File(s)	Summary
Candidate validation and selection experimental/softmax_reparameterization/search.py, tests/unit/torch/quantization/test_softmax_reparameterization.py, tests/gpu/torch/quantization/test_softmax_reparameterization_cuda.py	search_head validates inputs and quantization settings, independently quantizes shifted candidates, and selects by validation KL. Unit and CUDA tests cover candidate selection, source-model isolation, tie handling, invalid inputs, and save/restore behavior.
Usage and validation support experimental/softmax_reparameterization/__init__.py, experimental/softmax_reparameterization/example.py, experimental/softmax_reparameterization/README.md, experimental/softmax_reparameterization/requirements.txt, CHANGELOG.rst	Adds package exports and a synthetic example that compares the selected candidate with a zero-coefficient baseline. Documents usage, supported configurations, limitations, and test commands, and adds a changelog entry.

Priority: ⬇️ Low

Estimated code review effort: 3 (Moderate) | ~20 minutes

Change: Feature

Sequence Diagram(s)

sequenceDiagram
  participant Caller
  participant search_head
  participant ModelOpt
  participant head_kl
  Caller->>search_head: Head, fit states, validation states, and quantization config
  loop Each coefficient
    search_head->>ModelOpt: Calibrate an independent candidate with fit states
    search_head->>head_kl: Measure candidate KL on validation states
    head_kl-->>search_head: Validation KL
  end
  search_head-->>Caller: Selected candidate, coefficient, and KL curve

Loading

Merge Risk: ⚪ Minimal · up to e8bf8

No actionable issue is established for this opt-in feature. Mergeability remains subject to normal validation of its supported configurations.

🚥 Pre-merge checks | ✅ 5 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 37.50% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 8 functions across 5 files. (3 skipped: 3…	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly identifies the main change: adding experimental softmax reparameterization for output-head quantization.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns	✅ Passed	No listed security anti-pattern was introduced. The PR adds no torch.load(..., weights_only=False), numpy.load(..., allow_pickle=True), hardcoded trust_remote_code=True, Python eval()/exec()…

Full details: Docstring Coverage

Explanation

Docstring coverage is 37.50% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 8 functions across 5 files. (3 skipped: 3 unsupported.)

• Fix all pre-merge checks with AI

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create a new PR

• Autopilot · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Warning

Some tools did not complete. Review the errors below.

🔧 OSV Scanner (2.6.0)

experimental/softmax_reparameterization/requirements.txt

OSV Scanner exited with code 128 without a usable report

---

_{Comment @coderabbitai help to get the list of available commands.}