feat(quant): add constant_amax to pin activation input_scale (NVFP4 experts input_scale=1.0)#1947
feat(quant): add constant_amax to pin activation input_scale (NVFP4 experts input_scale=1.0)#1947Edwardf0t1 wants to merge 2 commits into
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (3)
🚧 Files skipped from review as they are similar to previous changes (3)
📝 WalkthroughWalkthroughAdds ChangesConstant amax feature
Estimated code review effort: 3 (Moderate) | ~25 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches📝 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
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
modelopt/torch/quantization/config.py (1)
665-696: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winForbid setting both constant-amax modes
modelopt/torch/quantization/config.py:665-696—use_constant_amaxandconstant_amaxcan both be set today. That makes fake-quant use FP8 E4M3 max while export still uses the pinned_amaxvalue, so the simulated and exported scales can silently diverge. Add amodel_validatorto reject the combination.🤖 Prompt for 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. In `@modelopt/torch/quantization/config.py` around lines 665 - 696, The quantization config currently allows both use_constant_amax and constant_amax to be set at the same time, which can make the fake-quant path and exported scaling diverge. Add a model-level validator on the config class in config.py that checks these two fields together and raises a clear validation error when both are provided, while keeping the existing field-level validate_constant_amax check for positivity.
🧹 Nitpick comments (1)
modelopt_recipes/general/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast.yaml (1)
46-61: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDuplicated
constant_amax/nvfp4block acrossexperts.*andblock_sparse_moe*patterns.Both quantizer-name blocks apply identical
cfgpayloads (nvfp4 import +constant_amax: 2688.0) just to cover two different MoE naming conventions. A YAML anchor could remove the duplication and prevent the two copies from drifting if the constant ever needs updating.♻️ Optional refactor using YAML anchors
quant_cfg: - $import: base_disable_all - - quantizer_name: '*.experts.*weight_quantizer' - cfg: - $import: nvfp4 + - quantizer_name: '*.experts.*weight_quantizer' + cfg: &nvfp4_weight + $import: nvfp4 - quantizer_name: '*.experts.*input_quantizer' - cfg: - $import: nvfp4 - # 2688 = E2M1_MAX * E4M3_MAX (6 * 448) -> exported input_scale == 1.0, no calibration. - constant_amax: 2688.0 + cfg: &nvfp4_input_pinned + $import: nvfp4 + # 2688 = E2M1_MAX * E4M3_MAX (6 * 448) -> exported input_scale == 1.0, no calibration. + constant_amax: 2688.0 - quantizer_name: '*block_sparse_moe*weight_quantizer' - cfg: - $import: nvfp4 + cfg: *nvfp4_weight - quantizer_name: '*block_sparse_moe*input_quantizer' - cfg: - $import: nvfp4 - # 2688 = E2M1_MAX * E4M3_MAX (6 * 448) -> exported input_scale == 1.0, no calibration. - constant_amax: 2688.0 + cfg: *nvfp4_input_pinned - $import: kv_fp8_cast - $import: default_disabled_quantizers🤖 Prompt for 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. In `@modelopt_recipes/general/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast.yaml` around lines 46 - 61, Refactor the duplicated MoE quantizer config by defining a shared YAML anchor for the repeated nvfp4 payload and constant_amax value, then reuse it for both the '*.experts.*input_quantizer' and '*block_sparse_moe*input_quantizer' entries (and the matching weight_quantizer blocks if appropriate) so the two naming conventions stay in sync. Locate the duplicated configuration in the ptq recipe section that sets $import: nvfp4 and constant_amax: 2688.0, and replace the repeated blocks with aliases while preserving the existing matching patterns and behavior.
🤖 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/_test_utils/torch/quantization/tensor_quantizer_common.py`:
- Around line 294-307: The test method `test_constant_amax_nvfp4_input_scale`
imports `NVFP4QTensor` locally without any stated reason, which violates the
test import convention here. Move `from
modelopt.torch.quantization.qtensor.nvfp4_tensor import NVFP4QTensor` to the
module-level imports in `tensor_quantizer_common.py` with the other test
dependencies, and remove the in-method import so the test uses the shared
top-level symbol.
---
Outside diff comments:
In `@modelopt/torch/quantization/config.py`:
- Around line 665-696: The quantization config currently allows both
use_constant_amax and constant_amax to be set at the same time, which can make
the fake-quant path and exported scaling diverge. Add a model-level validator on
the config class in config.py that checks these two fields together and raises a
clear validation error when both are provided, while keeping the existing
field-level validate_constant_amax check for positivity.
---
Nitpick comments:
In
`@modelopt_recipes/general/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast.yaml`:
- Around line 46-61: Refactor the duplicated MoE quantizer config by defining a
shared YAML anchor for the repeated nvfp4 payload and constant_amax value, then
reuse it for both the '*.experts.*input_quantizer' and
'*block_sparse_moe*input_quantizer' entries (and the matching weight_quantizer
blocks if appropriate) so the two naming conventions stay in sync. Locate the
duplicated configuration in the ptq recipe section that sets $import: nvfp4 and
constant_amax: 2688.0, and replace the repeated blocks with aliases while
preserving the existing matching patterns and behavior.
🪄 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: 6f80259d-94df-4eaa-94a2-3ec0bf4d4b80
📒 Files selected for processing (6)
CHANGELOG.rstmodelopt/torch/quantization/config.pymodelopt/torch/quantization/model_calib.pymodelopt/torch/quantization/nn/modules/tensor_quantizer.pymodelopt_recipes/general/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast.yamltests/_test_utils/torch/quantization/tensor_quantizer_common.py
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1947 +/- ##
==========================================
- Coverage 77.80% 74.61% -3.19%
==========================================
Files 519 519
Lines 57960 57977 +17
==========================================
- Hits 45095 43262 -1833
- Misses 12865 14715 +1850
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:
|
| if val is not None: | ||
| # Pin amax to the constant on the _amax buffer so it is used by both the | ||
| # fake-quant forward pass and export; calibration is skipped in model_calib. | ||
| self.amax = float(val) |
There was a problem hiding this comment.
Hi Zhiyu, the pinned _amax buffer is registered on CPU and never moved to the model's device, will this have any implications. the test does the .to(device) so no errors are caught
There was a problem hiding this comment.
Good catch — you're right. The buffer was registered on CPU at config time and would not follow a later model.to(device), so the first post-calibration GPU forward would hit the amax.is_cuda assert in the dynamic-block-quant kernel (the .to(device) in the test masked it). Fixed in f224624: _get_amax now aligns the buffer with the input device on the fly. Thanks!
Edwardf0t1
requested review from
chadvoegele,
meenchen,
rohansjoshi and
shengliangxu
| This differs from ``use_constant_amax``, which pins amax to the FP8 E4M3 range (448.0) | ||
| for KV-cache cast math and does not register an ``_amax`` buffer. For NVFP4 activations | ||
| the exported ``input_scale`` equals ``amax / (E2M1_MAX * E4M3_MAX) = amax / (6 * 448)``; | ||
| setting ``constant_amax`` to ``2688.0`` therefore yields an exported ``input_scale`` of |
There was a problem hiding this comment.
Can the old use_constant_amax behavior be replicated by setting constant_amax = 448? If so, could we unify these args?
There was a problem hiding this comment.
Not exactly today — they overlap in intent but differ in mechanism:
use_constant_amax=Truepins amax to 448 only in the forward path (_get_amaxreturns it on the fly) and deliberately registers no_amaxbuffer. It's used for the KV*[kv]_bmm_quantizer, whose scale is exported through the dedicatedkv_cache_scaling_factorpath, not the_amax→input_scalepath. It's also format-agnostic (always 448).constant_amax=448would register_amax=448as a real buffer, so it lands in the state_dict and gets exported as aninput_scale(= amax/maxbound). For a KV bmm quantizer that's an extra/unexpected buffer next tokv_cache_scaling_factor.
So constant_amax=448 replicates the fake-quant behavior but not the export/state behavior. Unifying is feasible as a follow-up, but it touches the KV-cache export path (regression risk), so I'd want to first confirm the KV cast recipes still export byte-identical checkpoints before collapsing the two flags — I'd prefer to keep that out of this PR. Happy to file a follow-up issue to unify (likely: make constant_amax able to skip buffer registration / forward-only mode, then have the KV recipe set constant_amax: 448). WDYT?
…xperts input_scale=1.0) Add a `constant_amax` field to QuantizerAttributeConfig that pins a quantizer's amax to a fixed value and skips activation calibration (no forward statistics collected). Unlike `use_constant_amax` (hardcoded to 448.0 for KV-cache cast math, registers no buffer), `constant_amax` stores the constant on the `_amax` buffer so it is used by both the fake-quant forward and the exported scaling factor. For NVFP4 activations, input_scale = amax / (E2M1_MAX * E4M3_MAX) = amax / (6 * 448), so constant_amax=2688.0 yields input_scale == 1.0. Ships recipe nvfp4_experts_only_input_scale1-kv_fp8_cast.yaml that pins the MoE expert activation quantizers to input_scale=1.0 with no activation calibration, and unit tests covering buffer registration, calibration skip, the NVFP4 input_scale=1.0 mapping, and the positive validator. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Zhiyu Cheng <zhiyuc@nvidia.com>
Address PR review feedback: - Fix device mismatch: the constant_amax _amax buffer is registered on CPU at config time and may not follow a later model.to(device); align it with the input device in _get_amax so the post-calibration GPU forward doesn't hit the dynamic-block-quant is_cuda assert (juhi). - Add a model_validator forbidding use_constant_amax + constant_amax together, which would otherwise diverge fake-quant (448.0) from the exported pinned amax (CodeRabbit). - Move the NVFP4QTensor import in the test to module level per the repo test-import convention (CodeRabbit). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Zhiyu Cheng <zhiyuc@nvidia.com>
Edwardf0t1
force-pushed
the
feat/nvfp4-constant-amax-input-scale
branch
from
f224624 to
4e293de
Compare
4 participants
What does this PR do?
Type of change: new feature
Adds a
constant_amaxfield toQuantizerAttributeConfigthat pins a quantizer'samaxto a fixed value and skips activation calibration (no forward statistics collected). Unlike the existinguse_constant_amax— which hardcodes the amax to the FP8 E4M3 range (448.0) for KV-cache cast math and registers no_amaxbuffer —constant_amaxstores the constant on the_amaxbuffer so it is used by both the fake-quant forward pass and the exported scaling factor (e.g.input_scale).Motivation: pin the NVFP4 routed-expert activation
input_scaleto exactly1.0in the exported checkpoint, with no activation calibration. For NVFP4 the exported activation scale isinput_scale = amax / (E2M1_MAX * E4M3_MAX) = amax / (6 * 448), soconstant_amax = 2688.0yieldsinput_scale == 1.0.Changes:
config.py: newconstant_amax: float | Nonefield + positive-value validator.tensor_quantizer.py: setter registers the constant on the_amaxbuffer.model_calib.py:enable_stats_collection/finish_stats_collectionexclude constant-amax quantizers from calibration (so the fixed amax is never overwritten).modelopt_recipes/general/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast.yamlapplying it to the MoE expert activation quantizers.Usage
Or via the shipped recipe:
# modelopt_recipes/general/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast.yaml python examples/hf_ptq/hf_ptq.py --recipe nvfp4_experts_only_input_scale1-kv_fp8_cast ...Testing
tests/_test_utils/torch/quantization/tensor_quantizer_common.py(run viatest_tensor_quantizer_cpu.py):test_constant_amax_registers_amax—_amaxbuffer pinned to the constant; default (unset) registers no buffer.test_constant_amax_nvfp4_input_scale—constant_amax=2688⇒ exported NVFP4input_scale == 1.0.test_constant_amax_skips_calibration— quantizer excluded from calibration; amax not overwritten.test_constant_amax_must_be_positive— validator rejects non-positive values.pytest tests/unit/torch/quantization/test_tensor_quantizer_cpu.py tests/unit/torch/quantization/test_config_validation.py— 109 passed.pre-commit run --files ...— all hooks pass (incl.validate modelopt recipes).modelopt.recipe.loader.load_recipe, with$import: nvfp4correctly mergingconstant_amax: 2688.0.Before your PR is "Ready for review"
None; no behavior change when unset)CONTRIBUTING.md: N/AAdditional Information
The exported per-block E4M3 activation scales remain dynamic; only the per-tensor global scale (
input_scale) is pinned to 1.0.🤖 Generated with Claude Code
Summary by CodeRabbit
amaxto a user-specified constant, keeping behavior consistent for export and inference.amaxquantizers are excluded during calibration and restored afterward.amaxbehavior, validation rules (positive value, mutual exclusivity), and NVFP4 scaling/export results.