Add nvfp4_act_headroom activation calibration for NVFP4 - #2028
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds the ChangesNVFP4 activation headroom calibration
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant QuantizeMode
participant nvfp4_act_headroom_calibrate
participant TensorQuantizer
participant max_calibrate
QuantizeMode->>nvfp4_act_headroom_calibrate: select calibration mode
nvfp4_act_headroom_calibrate->>TensorQuantizer: replace eligible input calibrators
nvfp4_act_headroom_calibrate->>max_calibrate: run calibration loop
max_calibrate->>TensorQuantizer: collect activation statistics
nvfp4_act_headroom_calibrate->>TensorQuantizer: restore original calibrators
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Comment |
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@modelopt/torch/quantization/model_calib.py`:
- Around line 494-533: Update _swap_in_nvfp4_act_headroom_calibrators to recurse
through SequentialQuantizer wrappers, matching the traversal used by the amax
sync helpers. Evaluate wrapped TensorQuantizer leaves as NVFP4 input quantizers
even when the outer name does not end with input_quantizer, then install
NVFP4ActHeadroomCalibrator and increment the swap count for each qualifying
leaf.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: aa736bbf-7a08-4774-93cb-5b3745245436
📒 Files selected for processing (9)
CHANGELOG.rstmodelopt/torch/quantization/calib/__init__.pymodelopt/torch/quantization/calib/nvfp4_act_headroom.pymodelopt/torch/quantization/config.pymodelopt/torch/quantization/mode.pymodelopt/torch/quantization/model_calib.pymodelopt_recipes/configs/ptq/units/a4_nvfp4.yamlmodelopt_recipes/general/ptq/nvfp4_act_only_headroom-kv_fp8_cast.yamltests/unit/torch/quantization/test_nvfp4_act_headroom.py
d88d34d to
c6853dc
Compare
NVFP4 scales a tensor with an FP8-E4M3 scale per 16-element block plus one per-tensor global scale. Plain max calibration sets that global scale from the largest block seen during calibration, leaving no room above it: any activation larger than the calibration max saturates. nvfp4_act_headroom instead anchors the global scale to a low percentile of the per-block amax distribution, amax = max(rho * anchor, floor), placing the calibrated blocks in the lower part of the FP8 scale range and leaving the rest as headroom. A high-percentile floor keeps the scale at or above the calibrated blocks so calibration data is never clipped, and a guardrail warns when the observed range is wider than FP8 scales can represent. Applies to NVFP4 dynamic-block input quantizers only; weight quantizers and all other quantizers are calibrated with plain max. - calib/nvfp4_act_headroom.py: NVFP4ActHeadroomCalibrator - config.py: NVFP4ActHeadroomCalibConfig (anchor_percentile, rho) - model_calib.py: nvfp4_act_headroom_calibrate - mode.py: NVFP4ActHeadroomCalibrateModeDescriptor - modelopt_recipes: a4_nvfp4 unit and nvfp4_act_only_headroom-kv_fp8_cast, an activation-only recipe whose module coverage matches nvfp4_default-kv_fp8_cast Pre-commit hooks were run manually on these files (all pass); the commit hook is skipped only because its autostash cannot write to a read-only cache in this environment. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
c6853dc to
14a74cb
Compare
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@CHANGELOG.rst`:
- Line 11: Update the nvfp4_act_headroom changelog entry to state that the
high-percentile floor is not guaranteed to prevent clipping for every
calibration sample, and that the guardrail evaluates floor/anchor rather than
the full observed amax range. Add the export limitation that activation-only
scales may be dropped when checkpoint format is inferred from weight quantizers,
while clarifying that the W4A4 recipe is unaffected because it includes NVFP4
weights.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: f1be1280-8521-4eeb-9642-b00eba4f41ca
📒 Files selected for processing (8)
CHANGELOG.rstmodelopt/torch/quantization/calib/__init__.pymodelopt/torch/quantization/calib/nvfp4_act_headroom.pymodelopt/torch/quantization/config.pymodelopt/torch/quantization/mode.pymodelopt/torch/quantization/model_calib.pymodelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yamltests/unit/torch/quantization/test_nvfp4_act_headroom.py
🚧 Files skipped from review as they are similar to previous changes (6)
- modelopt/torch/quantization/calib/init.py
- modelopt/torch/quantization/calib/nvfp4_act_headroom.py
- modelopt/torch/quantization/config.py
- modelopt/torch/quantization/mode.py
- tests/unit/torch/quantization/test_nvfp4_act_headroom.py
- modelopt/torch/quantization/model_calib.py
- calib/nvfp4_act_headroom.py: floor the result at the largest observed per-block amax instead of the 99.99th percentile. The percentile floor could land below the observed max (bin centers round down, and 0.01% of blocks sit above it by construction), so the guardrail path could return a scale that clips calibration data -- worse than plain max. Warn whenever the anchor term cannot clear the observed max (headroom collapses at rho, not at the FP8 window), zero pad a ragged last dim instead of tripping the divisibility assert in reduce_block_amax, and assert on NaN/Inf as MaxCalibrator does. - Reject anchor_percentile=0 in both the config and the calibrator: searchsorted returns bin 0 for a zero target regardless of the low-tail mask, pinning the anchor to 2**-40. - model_calib.py: restore the original calibrators after calibration so the algorithm cannot leak into a later max_calibrate on the same model, and warn when no NVFP4 activation quantizer matches. - Document the recipe in modelopt_recipes/ptq.md and add it to the builtin recipe list in tests/unit/recipe/test_loader.py (two recipe-doc tests were failing). - Tests: assert the no-clipping invariant strictly across six distributions including the guardrail regime, plus regressions for percentile 0, ragged last dim, NaN/Inf, and calibrator restore. Pre-commit hooks were run manually on these files (all pass); the commit hook is skipped only because its autostash cannot write to a read-only cache in this environment. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/unit/torch/quantization/test_nvfp4_act_headroom.py`:
- Line 143: Move the max_calibrate import from the test body to the module-level
imports in test_nvfp4_act_headroom.py, alongside the existing model_calib helper
imports, and remove the local import.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: f26391ac-762d-4c3d-9de3-8867ce7c6a97
📒 Files selected for processing (7)
CHANGELOG.rstmodelopt/torch/quantization/calib/nvfp4_act_headroom.pymodelopt/torch/quantization/config.pymodelopt/torch/quantization/model_calib.pymodelopt_recipes/ptq.mdtests/unit/recipe/test_loader.pytests/unit/torch/quantization/test_nvfp4_act_headroom.py
🚧 Files skipped from review as they are similar to previous changes (4)
- CHANGELOG.rst
- modelopt/torch/quantization/config.py
- modelopt/torch/quantization/model_calib.py
- modelopt/torch/quantization/calib/nvfp4_act_headroom.py
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #2028 +/- ##
==========================================
+ Coverage 66.83% 69.09% +2.26%
==========================================
Files 519 520 +1
Lines 58916 60378 +1462
==========================================
+ Hits 39376 41720 +2344
+ Misses 19540 18658 -882
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:
|
This algorithm delegates all weight calibration to max_calibrate, but did not accept the shared-state knobs its siblings expose, so switching a recipe from `max` to `nvfp4_act_headroom` silently lost the ability to override the weight global_amax grouping patterns or to share one weight amax across local experts in a SequentialMLP MoE layer. Defaults are unchanged: resolve_patterns(None) already returns the default patterns, so behavior only changes when a recipe sets these fields. Also document the data-parallel sync semantics: max_calibrate combines per-rank scales with a MAX all-reduce, so the result is the largest per-rank headroom scale rather than the scale implied by pooling every rank's distribution. Pre-commit hooks were run manually on these files (all pass); the commit hook is skipped only because its autostash cannot write to a read-only cache in this environment. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
Changing how a percentile is reconstructed from the histogram (bin center vs. a bin edge vs. cdf interpolation within the bin) shifts every calibrated scale, so record that the current choice is intentional rather than an oversight. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
meenchen
left a comment
meenchen
left a comment
There was a problem hiding this comment.
Bot review (gpt-5.6-sol) — DM the bot to share feedback.
The calibration approach and coverage look thoughtful, but I found one correctness issue in the advertised FP8-range guardrail and two required review concerns. Design review: the problem is to choose an NVFP4 activation global scale from the per-block distribution while retaining unseen-outlier headroom with bounded collection memory. Existing alternatives include extending/composing the in-repo HistogramCalibrator with reduce_block_amax, implementing this as max calibration plus a post-processing collector, or using PyTorch/NumPy/SciPy percentile primitives (all already available). The bounded log2 histogram is a plausible reason for a specialized collector, but the PR body does not explicitly compare it with the existing histogram calibrator/post-processing path; please document why extension/reuse is unsuitable. Licensing also needs human sign-off: the two new files use a 2024 NVIDIA header, while the repository's canonical LICENSE_HEADER currently says 2026, so the new headers do not qualify for the exact-header auto-approval exception. The PR is cohesive and under the size threshold, and its unit/integration tests otherwise cover the main behavior well.
| assert amax >= float(x.abs().max()) | ||
|
|
||
|
|
||
| @pytest.mark.parametrize("bad", [float("nan"), float("inf")]) |
There was a problem hiding this comment.
Bot comment.
Please move this import, and the model_calib as mc import in test_shared_state_knobs_are_forwarded_to_max_calibrate, to the import block at the top of the file. There is no circular/optional-dependency reason given for keeping these Python imports local.
There was a problem hiding this comment.
Done in f4ccecc — both this one and the model_calib import in test_shared_state_knobs_are_forwarded_to_max_calibrate are now in the module-level import block. There was no circular or optional-dependency reason for keeping them local.
| # The calibrators are restored afterwards so this algorithm does not leak into a later | ||
| # calibration of the same model, and so a repeat run starts from a fresh histogram. | ||
| try: | ||
| max_calibrate( |
There was a problem hiding this comment.
no. We deliberately focusing on max calib for weights here and only goal is to be able to generate the activation scales. Doing MSE will make the implementation more complex. Later on I hope we can find a way to split weight and activation calib algorithms separately.
There was a problem hiding this comment.
Good question. It composes cleanly — mse_calibrate itself starts with max_calibrate(...) and then refines only weights, so the swapped-in activation calibrators still collect during that forward pass and the MSE weight loop never touches them. I checked: substituting mse_calibrate here leaves the activation scale bit-identical (19274.088 either way).
Two caveats before we do it:
- For this recipe it would currently be a no-op on weights. With
fp8_scale_sweep=Truethe MSE weight path applies only to static NVFP4 weights (_uses_modelopt_fp8_weight_scalesrequiresis_nvfp4_static), and the shipped recipe uses dynamic-block NVFP4 weights — I measured the weight amax unchanged at0.088378318. It would needfp8_scale_sweep=False(the multiplier search) or static weights to do anything. - It should be a knob, not a hardcoded swap, since weight calibration is an orthogonal axis to the activation scale this algorithm sets.
Worth noting there is a real gap here: chaining does not work as a substitute. algorithm: [mse, nvfp4_act_headroom] has the second step's max_calibrate reset the MSE weights, and [nvfp4_act_headroom, mse] has the MSE step's internal max_calibrate reset the activation scales. So a knob on this call is currently the only way to get MSE weights and headroom activations in one pass.
Happy to add a weight-calibration option in this PR if you'd like it here, or keep this PR to the activation scale and do it as a follow-up — your call.
There was a problem hiding this comment.
Can we call _run_scale_calibration1
our scale algorithms are focused on weight scale setting (we should update the _run_scale_calibration to say _run_weight_scale_calibration to convey this cleanly). It is orthogonal to the activation scale setting targeted in this PR.
There was a problem hiding this comment.
RB: Wdyt? Do you think what I proposed above is a good idea?
There was a problem hiding this comment.
🤖 Bot comment.
🐝 Looking into the calibration call flow and state interactions now.
There was a problem hiding this comment.
🤖 Bot comment.
Yes — I think that is the right abstraction. The temporary calibrator swap defines the activation-scale policy; then a renamed _run_weight_scale_calibration can independently choose max / mse / local_hessian. Both MSE and local-Hessian start with max_calibrate, so the swapped headroom collector is populated in that initialization pass, while their later refinement only changes weights. It also fixes the ordering problem with chaining two top-level algorithms.
I’d add three guardrails:
- expose a typed nested
weight_scale_algorithmand default it explicitly to{"method": "max"}here, rather than inheriting the helper’s currentNone -> msedefault; - keep kwargs owned by that nested algorithm so
distributed_sync/shared_statesare not duplicated; - handle
sync_expert_weight_amaxdeliberately: it is currently max-only, and MSE/local-Hessian do not accept or forward it.
Then test at least max + MSE composition (and ideally local-Hessian): activation amax remains headroom-derived while the selected weight calibration takes effect. With that contract cleanup, I’d prefer your proposal over hard-wiring max_calibrate here.
There was a problem hiding this comment.
Implemented in a596c0e — thanks, this is a better shape than what I had.
To be clear about the change of position: my earlier "no" was to hard-wiring mse here, and that objection stands. Your proposal is a different thing — it routes weight calibration through the dispatch helper that already exists, which is exactly the "split weight and activation calib algorithms separately" I was hoping for later. So it lands now rather than later.
algorithm:
method: nvfp4_act_headroom
anchor_percentile: 1 # activation policy, owned here
upper_percentile: 99.99
rho: 16384
weight_scale_algorithm: # weight policy, delegated
method: max # or mse / local_hessian, plus that algorithm's options_run_scale_calibration is renamed _run_weight_scale_calibration as you suggested (private, and I updated the LSQ caller and its tests). distributed_sync, shared_states and sync_expert_weight_amax move off the top level into the nested config, which also resolves the third guardrail on its own: MaxCalibConfig defines sync_expert_weight_amax and MseCalibConfig/LocalHessianCalibConfig do not, so the union already types it per method — there was nothing to plumb, only a duplicate to delete.
Two implementation notes worth recording:
- The default must carry
methodexplicitly.{"method": "max"}rather thanNoneavoids the helper'sNone -> msedefault, as your bot flagged. It also can't be a bareMaxCalibConfig(): the helper doesmodel_dump(exclude_unset=True), so an unsetmethodis dropped andcalib_funcs[None]raisesKeyError. - The
field_serializeris load-bearing, not cosmetic.wrapped_calib_funcdoes a fullconfig.model_dump(), so without the sparse-shape serializer every nested default is forwarded as a kwarg andmax_calibrateraises onmoe_calib_experts_ratio/layerwise. Copying the LSQ pattern fixes it; verified the dispatch receives{}for the default and{"fp8_scale_sweep": True}for{"method": "mse", "fp8_scale_sweep": true}.
One caveat on the test you asked for. Asserting that the selected weight calibration takes effect isn't possible in the CPU unit environment: MSE on static NVFP4 needs triton, and on dynamic NVFP4 it needs CUDA (I hit static_blockwise_fp4_fake_quant requires triton and amax must be a CUDA tensor respectively). Worth knowing regardless of this PR: with fp8_scale_sweep=True the MSE weight path applies only to static NVFP4, so on this recipe's dynamic-block weights it is a no-op — a test asserting "weights changed" would pass vacuously. So the unit tests assert the dispatch and options plumbing plus survival of the headroom activation scale, and I've left the numerical effect to GPU coverage. Happy to add a GPU test if you'd prefer it in this PR.
31 tests on the calibrator; tests/unit/torch/quantization/ and tests/unit/recipe/ are back to the pre-existing baseline with no new failures.
…observed max An earlier commit made the observed per-block max the hard floor for the global scale, on the grounds that a high percentile does not guarantee every calibration block is representable. Measuring the resulting quantization error shows that was the wrong trade. Flooring at the literal max means a single freak block drags the global scale up until every other block's FP8 block scale falls below subnormal: on a tensor with one block seven orders of magnitude above the rest, it flushed 99.998% of elements to zero, versus 6.7% when the rare blocks are clipped instead. On a benign wide-range tensor the two differ by 0.4% relative MSE, so the strict no-clipping floor buys almost nothing where it is safe and is destructive where it is not. Restore the top of the calibrated range as a configurable percentile (``upper_percentile``, default 99.99, ``100`` selects the literal observed max), and state the trade-off in the docstrings rather than asserting a no-clipping guarantee the default does not provide. That overclaim, not the percentile, was the actual defect. This also makes the guardrail ratio self-consistent again: the reported range is measured against the bound that is actually operative. Also address review feedback: move test imports to module level, and use the canonical 2026 copyright year on the two new files. Pre-commit hooks were run manually on these files (all pass); the commit hook is skipped only because its autostash cannot write to a read-only cache in this environment. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
meenchen
left a comment
meenchen
left a comment
There was a problem hiding this comment.
LGTM implementation-wise. Curious to know:
- Do you see significant differences between
amaxandrho × anchor - With this calibration, do we expect to use the headromm act in most of the case, and the calibrated amax would be the fallback for rare cases?
Yes: Please see this file: https://github.com/NVIDIA/Model-Optimizer/blob/72dbfb9ed2a20ec901671708c29799f622a83426/experimental/nvfp4_global_scale_study/input_scale_comparison.html
So far I find this method to be very effective for nemotron models but not for Qwen3.6. I think we can give it a try for other OSS models and see how it compares with the default max algo. |
|
Thanks — all three points were valid. Addressed in f4ccecc. On the guardrail ( Measuring it showed the change was wrong. Flooring the global scale at the literal max means a single freak block drags the scale up until every other block's FP8 block scale falls below subnormal. On a tensor with one block seven orders of magnitude above the rest, that flushes 99.998% of elements to zero, against 6.7% when the rare blocks are clipped instead; on a benign wide-range tensor the two differ by 0.4% relative MSE. So the strict floor buys almost nothing where it is safe and is destructive where it is not. f4ccecc therefore restores the top of the calibrated range as a configurable percentile — This makes the guardrail self-consistent again rather than patched: Imports: moved Licensing: good catch — both new files now carry the canonical 2026 header. They had 2024 because I copied the header from Design justification: added to the PR description. Short version: Verification: the restored calibrator is bit-identical to the reference implementation at both |
The recipe only set anchor_percentile, leaving upper_percentile and rho implicit. Spell all three out at their defaults so the recipe documents the tunables and the formula they feed, and note that upper_percentile=100 selects the literal observed max. Values match the defaults, so calibration behavior is unchanged. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
meenchen
left a comment
meenchen
left a comment
There was a problem hiding this comment.
Bot review (gpt-5.6-sol) — DM the bot to share feedback.
Design review: this solves bounded-memory selection of an NVFP4 activation global scale from per-block amaxes while retaining outlier headroom. The main alternatives are composing/extending the existing HistogramCalibrator with reduce_block_amax, max calibration plus post-processing, or Torch/NumPy/SciPy percentile primitives (Torch, NumPy, and SciPy are already dependencies). The updated PR body now reasonably explains why the existing linear value histogram and those composition paths do not provide the required derived statistic and low-tail log resolution, so the design concern is addressed.
Re-review: the guardrail/documentation concern and local-import concern are addressed. However, the critical SequentialQuantizer issue remains in the current branch: leaf names gain a numeric suffix and fail the endswith("input_quantizer") check, so those activations silently use max calibration. The promised regression test is also absent. Additionally, one new recipe still has a 2024 header rather than the canonical 2026 LICENSE_HEADER; this requires human licensing sign-off (and appears inconsistent with the author's statement that both new-file headers were fixed).
| def _is_nvfp4_dynamic_input_quantizer(name: str, module: nn.Module) -> bool: | ||
| """True for an enabled NVFP4 (E2M1 + E4M3 scale) dynamic-block *input* quantizer. | ||
|
|
||
| These carry their per-tensor global scale in ``_amax`` and are calibrated like plain max |
There was a problem hiding this comment.
Bot comment.
The prior SequentialQuantizer correctness issue is still present in the current branch. For an input wrapper, named_modules() exposes leaves as names such as layer.input_quantizer.0, so this suffix check rejects every wrapped leaf and nvfp4_act_headroom silently falls back to max for them. Please traverse each parent input quantizer through _iter_leaf_quantizers (or otherwise preserve the parent's input role while testing each leaf), and add a regression test with an NVFP4 leaf inside SequentialQuantizer.
| @@ -0,0 +1,47 @@ | |||
| # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |||
There was a problem hiding this comment.
Bot comment.
This new file still uses a 2024 copyright header, while the repository's canonical LICENSE_HEADER says 2026. Thus it does not qualify for the exact-standard-header auto-approval exception and needs either alignment with the canonical header or explicit human licensing sign-off.
…se year A quant_cfg whose input_quantizer cfg is a list builds a SequentialQuantizer, whose leaves are submodules named ``...input_quantizer.<i>``. The selector matched on the leaf name ending in ``input_quantizer``, so neither the container (not a TensorQuantizer) nor its leaves (name has a numeric suffix) qualified, and those NVFP4 activations silently fell back to plain max. Match on the container name and walk the leaves via the existing ``_iter_leaf_quantizers`` helper, with a regression test asserting all four leaves of a two-layer model are calibrated and receive the headroom scale rather than a max scale. The new recipe also still carried a 2024 copyright header, copied from the unit file it was derived from; it now uses the canonical 2026 header like the other two new files. Pre-commit hooks were run manually on these files (all pass); the commit hook is skipped only because its autostash cannot write to a read-only cache in this environment. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
|
Both remaining points were valid and are fixed in 136e036. SequentialQuantizer. You're right, and I was wrong to scope this out. I had treated it as theoretical; it isn't. A perfectly ordinary so those NVFP4 activations silently got plain max, exactly as you describe. The selector now matches on the container name and walks the leaves through the existing For the record on the "promised" test: the earlier CodeRabbit thread on this was auto-annotated "✅ Addressed in commit 14a74cb", which was incorrect — that commit only changed the recipe. I then read the auto-resolution as the issue being moot and explicitly declined it as out of scope rather than promising a test. That was a bad call regardless, since it is reachable from a normal recipe; thanks for pushing back on it. Licensing. Also correct, and my earlier statement was wrong. This PR adds three files, not two — I fixed the two Python files and missed the recipe YAML, which had inherited a 2024 header from the unit file I derived it from. All three now carry the canonical 2026 header: Tests are at 27 for the calibrator; |
…eaves A SequentialQuantizer wraps several quantizers over one activation, while this algorithm derives a single per-tensor global scale, so the composition has no defined behavior here. Calibrating each leaf independently, as the previous commit did, papers over that rather than settling it. Reject it explicitly instead: if an input quantizer is a SequentialQuantizer wrapping an NVFP4 dynamic-block quantizer, raise NotImplementedError naming the module. Validation runs before any calibrator is swapped, so a rejected model is left untouched. SequentialQuantizers with no NVFP4 leaf are irrelevant to this algorithm and are ignored rather than rejected. Pre-commit hooks were run manually on these files (all pass); the commit hook is skipped only because its autostash cannot write to a read-only cache in this environment. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
|
Changed approach in 4620291: Reasoning: a raise NotImplementedError(
f"nvfp4_act_headroom does not support SequentialQuantizer activation quantizers, "
f"but {name!r} is one wrapping an NVFP4 dynamic-block quantizer. Use a single "
f"activation quantizer config for these modules, or calibrate with 'max'."
)Three details worth noting:
Tests cover all three: the rejection with its message, that the model is unchanged after a rejection, and that a non-NVFP4 sequential activation is ignored. 29 tests for the calibrator; |
Weight-scale calibration is orthogonal to the activation global scale this
algorithm sets, but it was hard-wired to max_calibrate, so a recipe could not
combine, say, MSE weights with headroom activations. Chaining is not a substitute:
[mse, nvfp4_act_headroom] has the second step's max_calibrate reset the MSE
weights, and [nvfp4_act_headroom, mse] has MSE's internal max_calibrate reset the
activation scales.
Route weight calibration through the existing dispatch helper, renamed
_run_weight_scale_calibration to say what it actually governs, and expose it as a
nested, typed weight_scale_algorithm on the config (max / mse / local_hessian plus
that algorithm's own options). Default explicitly to {"method": "max"} rather than
inheriting the helper's None -> mse default, so selecting this activation policy
never silently changes how weights are calibrated.
distributed_sync, shared_states and sync_expert_weight_amax move off the top level
into that nested config, where the union types them per method -- MaxCalibConfig
defines sync_expert_weight_amax and the others do not. A field_serializer keeps the
sparse public dict shape, which is load-bearing: wrapped_calib_func does a full
model_dump, so without it every nested default would be forwarded as a kwarg and
max_calibrate would raise on moe_calib_experts_ratio and layerwise.
Tests cover the max default with sparse dispatch, mse selection with its options,
and nested knob ownership. Asserting that MSE actually moves the weights needs
triton (static NVFP4) or CUDA (dynamic), so these assert the dispatch and that the
headroom activation scale survives.
Pre-commit hooks were run manually on these files (all pass); the commit hook is
skipped only because its autostash cannot write to a read-only cache in this
environment.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
3 participants
What does this PR do?
Type of change: new feature
Adds
nvfp4_act_headroom, a calibration algorithm for NVFP4 activation global scales.NVFP4 scales a tensor in two levels: an FP8-E4M3 scale per 16-element block, plus one per-tensor global scale. Plain
maxcalibration sets that global scale from the largest per-block amax seen during calibration, which leaves no room above it — any activation larger than the calibration max saturates.nvfp4_act_headroominstead anchors the global scale to a low percentile of the per-block amax distribution:anchoris the per-block amax atanchor_percentile;upperis the per-block amax atupper_percentileand is the top of the range the scale commits to representing. Multiplying the low anchor byrhoplaces the calibrated blocks in the lower part of the FP8 scale range and leaves the rest as headroom.upper_percentiledefaults to 99.99 rather than the literal maximum on purpose. Flooring at the literal max means one freak block drags the global scale up until every other block's FP8 block scale falls below subnormal: on a tensor with a single block seven orders of magnitude above the rest, that flushes 99.998% of elements to zero, versus 6.7% when the rare blocks are clipped instead. On a benign wide-range tensor the two choices differ by 0.4% relative MSE. Setupper_percentile=100to floor at the literal observed max, which guarantees no calibration data is clipped, at that exposure. The calibrator warns when the per-block range is too wide forrhoto clear any headroom.The algorithm applies only to NVFP4 dynamic-block input quantizers; weight quantizers and everything else keep plain
max.Files:
calib/nvfp4_act_headroom.py—NVFP4ActHeadroomCalibrator(log2 histogram of per-block amaxes, bounded memory)config.py—NVFP4ActHeadroomCalibConfig(anchor_percentiledefault 1,rhodefault 16384)model_calib.py—nvfp4_act_headroom_calibratemode.py— mode registrationmodelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml— mirrorsnvfp4_default-kv_fp8_cast(dynamic NVFP4 W4A4 + FP8 KV-cache cast, same module coverage) with only the calibration algorithm swapped:Why a dedicated collector rather than reusing
HistogramCalibratorHistogramCalibratorhistograms the tensor values themselves into linear, dynamically re-ranged bins, and its percentile mode returns an amax that clips a high-tail fraction. This algorithm needs a different statistic (per-block amaxes, a derived quantity), different bin spacing (log2, because block amaxes span many decades and the anchor is read from the low tail, where linear bins have almost no resolution), and a different reduction (a low-percentile anchor scaled up, not an upper-tail clip). Reusing it would mean changing its collection semantics for every existing caller; composing it would still leave the log-spacing and the derived statistic unaddressed. The collector here is a fixed 512-bin int64 histogram per quantizer, so the memory argument for a histogram (rather than retaining values) is preserved.Usage
Or directly:
Testing
Unit tests —
tests/unit/torch/quantization/test_nvfp4_act_headroom.py, 26 tests covering the anchor and range terms, monotonicity inanchor_percentile, headroom above plain max, the too-wide-range fallback and its warning, the rare-outlier case (the default clips it;upper_percentile=100chases it), theupper_percentile=100no-clipping guarantee across six distributions, input validation (rho/ percentile bounds, NaN/Inf, all-zero, a last dim that is not a multiple of the block size), quantizer selection, calibrator restoration after calibration,shared_states/sync_expert_weight_amaxforwarding, and W4A4 end-to-end confirming weights fall back to plain max while only activations get the headroom scale.Full
tests/unit/torch/quantization/andtests/unit/recipe/suites pass with no regressions (the remaining failures in this environment are a pre-existing broken-torchvisionimport andtest_data_parallel_auto_quantize, both of which fail identically onmain).End-to-end on a real model — ran the shipped recipe on a 30B hybrid Mamba/attention MoE (1024 calibration samples @ seq 4096) and verified the export is a standard NVFP4 checkpoint: 6268 activation quantizers calibrated; 19G vs the 62G BF16 source; 6260
input_scale, 6267weight_scale/weight_scale_2;quant_algo: NVFP4,kv_cache_quant_algo: FP8; 53 excluded modules; no module carrying aninput_scalewithout aweight_scale. 32 of 6268 quantizers (0.5%) had a per-block range too wide forrhoto clear and warned.Before your PR is "Ready for review"
CONTRIBUTING.md: N/A — no copied code, no new dependencies.🤖 Generated with Claude Code