feat(quant): skip max-calib forward when no activation needs data#1963
feat(quant): skip max-calib forward when no activation needs data#1963Edwardf0t1 wants to merge 3 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)
📝 WalkthroughWalkthroughMax calibration gains configurable detection of activation-statistics requirements. When none are needed, the activation ChangesMax calibration forward control
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant max_calibrate
participant TensorQuantizer
participant forward_loop
participant WeightCalibration
max_calibrate->>TensorQuantizer: Inspect enabled activation quantizers
TensorQuantizer-->>max_calibrate: Report calibration requirement
alt Activation statistics required
max_calibrate->>forward_loop: Run calibration forward loop
else No activation statistics required and skip enabled
max_calibrate->>WeightCalibration: Continue without forward loop
end
max_calibrate->>WeightCalibration: Calibrate weight tensors
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1963 +/- ##
==========================================
- Coverage 76.40% 76.28% -0.12%
==========================================
Files 522 522
Lines 58413 58821 +408
==========================================
+ Hits 44632 44874 +242
- Misses 13781 13947 +166
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:
|
Edwardf0t1
force-pushed
the
feat/skip-forward-loop-constant-amax
branch
from
be9b904 to
977b74c
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
|
| or getattr(module, "_constant_amax", None) is not None | ||
| ): | ||
| continue | ||
| if getattr(module, "_calibrator", None) is not None: |
There was a problem hiding this comment.
Can we also check if the model is mx formats like MXFP4? if the activation quantizer is MXFP4/MXFP8 we can skip calibration on that as well.
There was a problem hiding this comment.
RB: Can you give us a plan on how to do this?
There was a problem hiding this comment.
RB: Never mind, I think module._dynamic captures that. Is that correct?
There was a problem hiding this comment.
Good question — actually module._dynamic does not capture MXFP4/MXFP8. _dynamic is only set from a top-level type: dynamic, but the MX numerics configs set type: dynamic inside block_sizes (block-level E8M0 scales), leaving the top-level _dynamic == False:
# configs/numerics/mxfp8.yaml
num_bits: e4m3
block_sizes: {-1: 32, type: dynamic, scale_bits: e8m0}So an MX activation quantizer would currently fall through to the _calibrator is not None check and be treated as needing the forward. MX formats carry no per-tensor amax at all (is_mx_format makes amax always None), so they're safe to skip. I'll add module.is_mx_format to the skip conditions in _needs_activation_forward_for_max_calib so MXFP4/MXFP8 are covered too. Thanks for catching this!
There was a problem hiding this comment.
🤖 Bot comment.
🐝 Looking into this.
There was a problem hiding this comment.
🤖 Bot comment.
Edward’s answer is correct. Concretely, the plan is:
- Keep
_dynamicfor top-level dynamic quantizers. - Add
module.is_mx_formatas a separate skip condition, because MX’stype: dynamiclives insideblock_sizesand leaves_dynamic == False. - Add a regression test proving an MX quantizer has
is_mx_format == True,_dynamic == False, and makes_needs_activation_forward_for_max_calib()returnFalse; retain the normal data-driven case returningTrue.
This is already implemented at 3f9d4e28, and the required unit/GPU/regression checks are green.
realAsma
left a comment
realAsma
left a comment
There was a problem hiding this comment.
Looks great!! Thanks for this PR
Address PR #1963 review feedback: - Make MaxCalibConfig.skip_forward_without_activation_calib opt-in (default False). Defaulting it True silently changed calibration for every max-PTQ user and broke weight-only quantization under DeepSpeed ZeRO-3 (test_nested_deepspeed_backward[3-quant_cfg0]): the calibration forward materializes sharded params, so skipping it left the ZeRO-3 param shards out of sync with the quantized weights. Enable it explicitly in the nvfp4_experts_only_input_scale1-kv_fp8_cast recipe. - Also skip the forward for MX formats (MXFP4/MXFP8) in _needs_activation_forward_for_max_calib (realAsma). MX configs set type: dynamic inside block_sizes (E8M0 per-block scales) but leave the top-level _dynamic False, so the _dynamic check did not catch them; they carry no per-tensor amax (is_mx_format -> amax is None), so add an explicit is_mx_format check. Update tests (opt-in default, MX/dynamic coverage) and changelog. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Zhiyu Cheng <zhiyuc@nvidia.com>
|
Pushed
Tests updated (opt-in default + MX/dynamic coverage), changelog updated, pre-commit + recipe validation green. CI re-running. |
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>
Address PR #1963 review feedback: - Make MaxCalibConfig.skip_forward_without_activation_calib opt-in (default False). Defaulting it True silently changed calibration for every max-PTQ user and broke weight-only quantization under DeepSpeed ZeRO-3 (test_nested_deepspeed_backward[3-quant_cfg0]): the calibration forward materializes sharded params, so skipping it left the ZeRO-3 param shards out of sync with the quantized weights. Enable it explicitly in the nvfp4_experts_only_input_scale1-kv_fp8_cast recipe. - Also skip the forward for MX formats (MXFP4/MXFP8) in _needs_activation_forward_for_max_calib (realAsma). MX configs set type: dynamic inside block_sizes (E8M0 per-block scales) but leave the top-level _dynamic False, so the _dynamic check did not catch them; they carry no per-tensor amax (is_mx_format -> amax is None), so add an explicit is_mx_format check. Update tests (opt-in default, MX/dynamic coverage) and changelog. 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/skip-forward-loop-constant-amax
branch
from
33cebce to
3f9d4e2
Compare
| :func:`weight_only_quantize`, so they never need the data forward. Activation-side | ||
| quantizers (input/output/BMM) need it only when they collect data-driven statistics: | ||
| i.e. they are enabled, have a calibrator, and are none of | ||
|
|
There was a problem hiding this comment.
should we also check for the bias_calibrator?
There was a problem hiding this comment.
Good catch — yes, we should. A quantizer with a static bias_calibrator collects data during the forward (finish_stats_collection → load_calib_bias), so it needs the forward even when its amax path is dynamic/MX (which carry no per-tensor amax). My original helper would have wrongly skipped that case.
Fixed in dec014d: _needs_activation_forward_for_max_calib now mirrors finish_stats_collection exactly — a non-constant quantizer with a static bias forces the forward:
if not is_constant and module.bias_calibrator is not None and module.bias_type == "static":
return TrueConstant-amax quantizers stay exempt (they continue before the bias block in finish_stats_collection, so their bias isn't calibrated either), and a dynamic bias is computed on the fly so still needs no forward. Added tests for MX+static-bias (→ forward runs), MX+dynamic-bias, and constant+static-bias. Thanks!
Address review feedback (juhi10071998): a quantizer with a static bias_calibrator collects data during the calibration forward (finish_stats_collection -> load_calib_bias), so it needs the forward even when its amax is dynamic/MX (which carry no per-tensor amax). _needs_activation_forward_for_max_calib now mirrors finish_stats_collection exactly: a non-constant quantizer with a static bias_calibrator forces the forward, while constant-amax quantizers (which are exempted from bias calibration) and dynamic-bias quantizers still do not. Add tests for MX + static/dynamic bias and constant + static bias. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Zhiyu Cheng <zhiyuc@nvidia.com>
3 participants
What does this PR do?
Type of change: new feature (performance)
Addresses @realAsma's review comment on #1947: when all enabled input/KV activation quantizers use a constant amax, max calibration still runs the full
forward_loop, which is wasted — the constant/dynamic quantizers collect no data-driven statistics.This adds
_needs_activation_forward_for_max_calib(model)and gates theforward_loopinmax_calibrate():use_constant_amax, orconstant_amax. Weight quantizers are still calibrated on the weight tensors directly viaweight_only_quantize(), so results are unchanged; only the wasted forward is avoided.Safety / scoping. The gate is behind an opt-in flag
max_calibrate(..., skip_forward_without_activation_calib=False).max_calibratehas many internal callers — the advanced algorithms that always need activations (MSE, local Hessian, SmoothQuant, SVDQuant, GPTQ) call it directly and keep the defaultFalse, so they are unaffected. Only the top-levelmaxpath opts in, viaMaxCalibConfig.skip_forward_without_activation_calib(defaultTrue).Also simplifies a default-arg lambda in
core_utils.sync_moe_expert_amax(lambda m, w=weight: m(w)→lambda m: m(weight)) so mypy can infer its type — surfaced by themax_calibratesignature change; theforward_loopis invoked synchronously so closure capture is safe.Usage
Automatic for the top-level
maxpath — no user action needed. For thenvfp4_experts_only_input_scale1-kv_fp8_castrecipe (all expert inputs + KV pinned to constant amax), calibration now performs weight-only calibration and skips the model forward:To force the old behavior: set
skip_forward_without_activation_calib: falseon the max algorithm config.Testing
New CPU unit tests in
tests/unit/torch/quantization/test_calib.py:test_max_calib_skips_forward_when_all_activations_constant— forward not called; weight quantizer still gets_amax; constant input amax preserved.test_max_calib_runs_forward_when_activation_needs_data— a normal input quantizer still triggers the forward.test_max_calib_default_does_not_skip_forward— direct callers (default flag) keep running the forward.test_needs_activation_forward_ignores_disabled_and_weight_quantizers.test_max_calib_config_enables_skip_by_default.pytest tests/unit/torch/quantization/test_calib.py→ 17 passed;test_quantize_cpu.py+test_tensor_quantizer_cpu.py→ 79 passed.pre-commiton all changed files passes (ruff, mypy, bandit, …).Before your PR is "Ready for review"
CONTRIBUTING.md: N/AAdditional Information
Follow-up to #1947 (now merged) per reviewer request. Closes the optimization discussion at #1947 (comment).
🤖 Generated with Claude Code
Summary by CodeRabbit