feat(quant): add constant_amax to pin activation input_scale (NVFP4 experts input_scale=1.0)#1947
Merged
Merged
Conversation
Contributor
|
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 |
Contributor
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
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% 77.18% -0.62%
==========================================
Files 519 519
Lines 57960 57977 +17
==========================================
- Hits 45095 44750 -345
- Misses 12865 13227 +362
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:
|
juhi10071998
reviewed
Jul 8, 2026
| 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) |
Contributor
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
Contributor
Author
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
chadvoegele
approved these changes
Jul 9, 2026
Comment on lines
+682
to
+685
| 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 |
Contributor
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?
Contributor
Author
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?
realAsma
reviewed
Jul 9, 2026
…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
Contributor
|
Do we still run full forward through the model even if all input/k/v quantizers use constant_amax? Can we check if all input/k/v quantizers are set to constant_amax? In that case can we override the forward_loop to do weight only calibration instead? |
Contributor
|
RB: Could you please provide a design plan for this? |
Contributor
🐝 Acknowledged. I’m checking the PR context and will reply with a design plan for #1947 (comment) shortly. |
Contributor
Design plan for #1947 (comment) / #1947 (comment): Yes, with the current code we still run the user Proposed implementation:
|
Contributor
Author
|
@realAsma Implemented this optimization in a stacked follow-up: #1963. It adds I scoped it opt-in at the |
Edwardf0t1
added a commit
that referenced
this pull request
Jul 11, 2026
When all enabled activation quantizers use a constant amax (constant_amax / use_constant_amax) or dynamic quantization, max calibration's forward pass collects no data-driven statistics and is pure overhead. Add _needs_activation_forward_for_max_calib() and gate the forward_loop in max_calibrate() behind a new opt-in flag skip_forward_without_activation_calib. Weight quantizers are still calibrated on the weight tensors directly via weight_only_quantize(), so results are unchanged; only the wasted forward is avoided (e.g. the nvfp4_experts_only_input_scale1 recipe). The flag defaults to False in max_calibrate() so the advanced algorithms that call it directly and always need activations (MSE, local Hessian, SmoothQuant, SVDQuant, GPTQ) are unaffected; it is enabled via MaxCalibConfig.skip_forward_without_activation_calib (default True) so only the top-level max path opts in. Also simplify a default-arg lambda in core_utils.sync_moe_expert_amax so mypy can infer its type (surfaced by the max_calibrate signature change); the forward_loop is invoked synchronously so closure capture is safe. Stacked on #1947 (depends on the constant_amax attribute). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Zhiyu Cheng <zhiyuc@nvidia.com>
cjluo-nv
added a commit
that referenced
this pull request
Jul 14, 2026
…tests (#1970) ### What does this PR do? **Type of change:** documentation (+ new tests) `modelopt_recipes/ptq.md` is meant to track every shipped PTQ recipe, but recent recipe PRs did not update it. This PR brings the doc back in sync and adds unit tests so it can't drift again. **Doc sync — general recipes:** - Add `nvfp4_experts_only_input_scale1-kv_fp8_cast` (#1947) to the shipped-recipes table (now 20 recipes) and document the `input_scale1` variant (constant amax 2688 → exported NVFP4 `input_scale == 1.0`, expert activations uncalibrated). - Add a pointer that the PTQ recipes' `quantize` sections also drive QAT/QAD, with the LSQ / Dual-LSQ QAD recipes under `general/qad/` (#1884). **Doc sync — model-specific recipes:** - Add `gemma4` (algorithm override, #1690), `diffusion_gemma` (extra `*self_conditioning*` exclusion, #1707), `vit` (FP8 with attention BMM quantizers for Torch-TRT, #1569), and the qwen3_5 / qwen3_5_moe `w4a16_nvfp4_mse-fp8_attn-kv_fp8_cast` twins (#1620). - Rewrite the checkpoint-mirror section for the `huggingface/models/nvidia/` tier: Super recipe rename (`super-nvfp4.yaml` → `nvfp4-mse.yaml` / `nvfp4-max-calib.yaml`), Nano-4B GGUF Q4_K_M mirror (#1327/#1606), Ultra-550B `nvfp4-4o6` (#1684); fix stale `models/<checkpoint>` paths. **Enforcement — new `tests/unit/recipe/test_recipe_docs.py`:** 1. Every `general/ptq/*.yaml` stem must appear (backticked) in `ptq.md`. 2. Every recipe row in the shipped-recipes table must exist on disk (catches renames/removals). 3. The "All N `general/ptq/` recipes" count must match the file count. 4. Every model folder under `huggingface/` containing `ptq/*.yaml` must be mentioned in the doc. ### Usage ```bash pytest tests/unit/recipe/test_recipe_docs.py -q ``` ### Testing - `pytest tests/unit/recipe/ -q` → 218 passed (includes the 4 new tests). - Verified enforcement bites: against the pre-update `ptq.md`, 3 of the 4 new tests fail with actionable messages. - `pre-commit run --files modelopt_recipes/ptq.md tests/unit/recipe/test_recipe_docs.py` → all hooks pass. ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: ✅ (docs + tests only; no runtime code changed) - 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?: ✅ — `tests/unit/recipe/test_recipe_docs.py` - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A (documentation/test change) - Did you get Claude approval on this PR?: ❌ (pending) ### Additional Information Follow-up to the recipe guide added in #1662. The doc-consistency tests could also be wired into `tools/precommit/check_modelopt_recipes.py` for commit-time feedback if reviewers prefer both layers. 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Documentation** * Added the `nvfp4_experts_only_input_scale1-kv_fp8_cast` recipe to the general PTQ catalog. * Documented the new `input_scale1` calibration behavior and related quantization/caching details. * Expanded and reworked PTQ documentation, including updated navigation, “Beyond PTQ” guidance, new/updated model-specific recipes, and expanded checkpoint mirror entries. * Updated the general PTQ shipped-recipe count to 20. * **Tests** * Added automated consistency checks between on-disk shipped recipe YAML files and the PTQ documentation (coverage, existence, and counts). <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Chenjie Luo <chenjiel@nvidia.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Edwardf0t1
added a commit
that referenced
this pull request
Jul 14, 2026
When all enabled activation quantizers use a constant amax (constant_amax / use_constant_amax) or dynamic quantization, max calibration's forward pass collects no data-driven statistics and is pure overhead. Add _needs_activation_forward_for_max_calib() and gate the forward_loop in max_calibrate() behind a new opt-in flag skip_forward_without_activation_calib. Weight quantizers are still calibrated on the weight tensors directly via weight_only_quantize(), so results are unchanged; only the wasted forward is avoided (e.g. the nvfp4_experts_only_input_scale1 recipe). The flag defaults to False in max_calibrate() so the advanced algorithms that call it directly and always need activations (MSE, local Hessian, SmoothQuant, SVDQuant, GPTQ) are unaffected; it is enabled via MaxCalibConfig.skip_forward_without_activation_calib (default True) so only the top-level max path opts in. Also simplify a default-arg lambda in core_utils.sync_moe_expert_amax so mypy can infer its type (surfaced by the max_calibrate signature change); the forward_loop is invoked synchronously so closure capture is safe. Stacked on #1947 (depends on the constant_amax attribute). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Zhiyu Cheng <zhiyuc@nvidia.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.