Add activation-aware local Hessian scale calibration#1976
Conversation
Extend local_hessian calibration with an opt-in act_quant_aware mode that minimizes the exact second-order expansion of ||XqWq - XW0||^2: Hessian built from quantized activations plus a cross term P = Xq^T(Xq - X)/B threaded through the NVFP4 MSE calibrator and the Triton FP8 sweep fast path. Default off is bit-identical to the current objective. Recommended usage for both modes: batch_size=1 and layerwise=True. Signed-off-by: realAsma <akuriparambi@nvidia.com>
|
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. |
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1976 +/- ##
=======================================
Coverage 76.33% 76.33%
=======================================
Files 522 522
Lines 58452 58505 +53
=======================================
+ Hits 44618 44662 +44
- Misses 13834 13843 +9
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
What does this PR do?
Type of change: new feature
Extends
local_hessianweight-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 / BP = X_qᵀ (X_q - X) / BThe default remains
False, preserving the existinglocal_hessianobjective. 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, usebatch_size=1andlayerwise=Truefor calibration fidelity and bounded memory.Usage
The new
nvfp4_local_hessian_act_awarerecipe alias enables this configuration directly.Testing
tests/unit/torch/quantization/test_local_hessian.py: 16 passedtests/unit/recipe/test_presets.py: 6 passedBefore your PR is "Ready for review"
CONTRIBUTING.md: N/AAdditional Information
The implementation follows the activation-quantization-aware block-wise output MSE derivation supplied with the feature request.