Add activation-aware local Hessian scale calibration

[realAsma]

What does this PR do?

Type of change: new feature

Extends local_hessian weight-scale calibration with an opt-in activation-quantization-aware block-wise output MSE objective.

When act_quant_aware=True, scale search minimizes the exact expanded error for ||X_q W_q - X W_0||² using:

• H = X_qᵀ X_q / B
• P = X_qᵀ (X_q - X) / B
• the quadratic Hessian term plus the activation-error cross term

The default remains False, preserving the existing local_hessian objective. The reference path and Triton FP8 scale-sweep fast path both support the cross term. Dense and fused-MoE activation capture are covered, with explicit fallback for transformed input quantizers.

For both regular and activation-aware local_hessian, use batch_size=1 and layerwise=True for calibration fidelity and bounded memory.

Usage

quant_cfg["algorithm"] = {
    "method": "local_hessian",
    "act_quant_aware": True,
    "fp8_scale_sweep": True,
}

The new nvfp4_local_hessian_act_aware recipe alias enables this configuration directly.

Testing

• tests/unit/torch/quantization/test_local_hessian.py: 16 passed
• Fused-MoE local-Hessian focused cases: 2 passed
• tests/unit/recipe/test_presets.py: 6 passed
• Activation-aware Triton/reference parity, validation, and end-to-end GPU cases: 5 passed
• Existing Hessian Triton/reference parity matrix: 18 passed
• Pre-commit on all changed files: passed

Before your PR is "Ready for review"

• Is this change backward compatible?: ✅
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A
• Did you write any new necessary tests?: ✅
• Did you update Changelog?: ❌ — pending while this PR is a draft
• Did you get Claude approval on this PR?: ❌ — pending while this PR is a draft

Additional Information

The implementation follows the activation-quantization-aware block-wise output MSE derivation supplied with the feature request.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[coderabbitai]

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 0ace0a14-af6a-463c-b4dc-9d0b7e7e5f4b

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch asma/activation-quant-aware-scale-setting

---

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

[github-actions]

PR Preview Action v1.8.1
🚀 View preview at https://NVIDIA.github.io/Model-Optimizer/pr-preview/pr-1976/
Built to branch gh-pages at 2026-07-18 20:17 UTC. Preview will be ready when the GitHub Pages deployment is complete.

[codecov]

Codecov Report

❌ Patch coverage is 74.73684% with 24 lines in your changes missing coverage. Please review.
✅ Project coverage is 77.78%. Comparing base (9392dfe) to head (3de41a5).

Files with missing lines	Patch %	Lines
...torch/kernels/quantization/gemm/nvfp4_fp8_sweep.py	0.00%	9 Missing ⚠️
modelopt/torch/quantization/model_calib.py	89.47%	8 Missing ⚠️
modelopt/torch/quantization/calib/mse.py	12.50%	7 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1976      +/-   ##
==========================================
- Coverage   77.79%   77.78%   -0.02%     
==========================================
  Files         518      518              
  Lines       58204    58279      +75     
==========================================
+ Hits        45279    45331      +52     
- Misses      12925    12948      +23     

Flag	Coverage Δ
unit	54.81% <74.73%> (+0.01%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[realAsma]

/claude review

[claude]

Claude review passed — no blocking issues found. LGTM

Findings: CRITICAL: 0, IMPORTANT: 0, SUGGESTION: 1

Scope

Full review. 12 files changed (+438/-43); reviewed all modelopt/ source (model_calib.py, calib/mse.py, config.py, kernels/.../nvfp4_fp8_sweep.py), the new recipe YAML, the README, and cross-checked mode.py registration and the test suite. Skipped only the CHANGELOG line and deprioritized test files (spot-checked for signature drift).

Assessment

This is a clean, well-guarded feature extension. Traced the activation-error-coupling math end-to-end and it holds:

• Sign flip dw = quant(w) - w (was w - quant(w)): the quadratic term is sign-invariant so the default path is numerically unchanged, while the new linear coupling term 2 dw^T (P W0) gets the correct sign. Applied consistently in the Triton kernel, the reference local_hessian_error, and docstrings.
• Coupling accumulation H = Xq^T Xq / B, P = Xq^T (Xq - X) / B, and the nbd,cnd->cnb einsum for P·W0 all match the block-wise expansion of the output MSE. The dropped scale-independent constant (hence possibly-negative objective) is documented, and argmin / best_loss=inf comparisons handle negative losses correctly.
• Kernel plumbing (flat coupling_bias layout, HAS_COUPLING constexpr, CUDA/numel validation) is consistent with x_flat; the non-coupling branch reuses x_flat as a harmless placeholder.
• Fallback logic for disabled/absent input quantizers and pre_quant_scale/rotation correctly reverts to the Hessian-only pre-hook with a rank-0 warning; the fused-MoE post-hook mirrors the dense post-hook.
• Backward compatibility: activation_error_coupling defaults False; internal renames (hessian->metric, hessian_for->metric_for) are private and all call sites/tests updated. Config flows cleanly via model_dump(). Post-calibration cleanup clears _coupling alongside _hessian, so no state leaks across runs.

Risk level: low. The one suggestion is a docstring accuracy nit on _LocalHessianAccumulator (still says H = sum X^T X, which is only exact on the default path). Non-blocking.

[realAsma]

/claude review

[claude]

Claude re-review — commit 1b30839 ("Use quantized activations for local Hessian capture")

Findings: CRITICAL: 0, IMPORTANT: 1, SUGGESTION: 0

Scope: Re-review of the follow-up commit on top of the previously-approved 2010ef8. Traced the delta end-to-end across model_calib.py, calib/mse.py, nvfp4_fp8_sweep.py (Triton kernel), config.py, the new preset YAML, and the test suite.

Assessment — the activation-error-coupling math is correct and well-tested:

• accumulate builds H = SUM X_q^T X_q and P = SUM X_q^T (X_q - X); the einsum for PW0 and the 2dW^T(PW0) term match the block-wise expansion of ||X_qW_q - X*W_0||^2. test_activation_error_coupling_objective_matches_output_mse verifies this against the true output MSE, including the deliberately-negative objective.
• The Triton kernel and reference local_hessian_error are mutually consistent: dw = quant(w)-w, flat coupling_bias layout matches x_flat, HAS_COUPLING gating, CUDA/numel validation, and best_loss=inf/argmin all handle negative losses correctly.
• Hook rework (input-quantizer register_forward_hook capturing X_q plus raw args[0]) is sound; is_quantized_linear guarantees input_quantizer exists, so the new continue never drops a real linear.
• Disabled/absent quantizer and pre_quant_scale/rotation fallback reverts to the Hessian-only path with a rank-0 warning; cleanup clears _coupling alongside _hessian.

Most impactful finding — [IMPORTANT Compatibility]: This commit changes the default (activation_error_coupling=False) Hessian from SUM X^T X to SUM X_q^T X_q whenever the input quantizer is enabled — which the released nvfp4_local_hessian recipe does. The existing recipe will yield different weight amaxes in 0.46 even without the new flag, yet the CHANGELOG says the change preserves the existing objective by default. Either document the default-path shift or gate X_q capture behind the new flag. Details inline on CHANGELOG.rst.

Risk level: low-moderate. The feature is clean and correct; the concern is the undocumented behavior shift to an already-shipped recipe.