Repository navigation
Conversation
Non-reentrant activation checkpointing ends its recompute early by raising inside the module forward. replace_function restored the original only on a normal exit, so a fused-experts module's F.linear interception stayed installed for the whole process and the next layer's recompute ran through it. With FSDP2 and activation checkpointing (the examples/llm_qat FSDP2 config) on Qwen3-MoE this crashes with "aten.mm.default got mixed torch.Tensor and DTensor". Restore the original in a finally. The same fix is part of NVIDIA#2745 (67d8ea0). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Shiyang Chen <shiychen@nvidia.com>
On the eager fused-experts path each expert's weight went through its own TensorQuantizer: one Triton launch per expert and projection, plus a weight[idx] slice whose backward writes a gradient the size of all experts once per expert. fp4_fake_quant_block now takes one global amax per index of dim 0, and StackedNVFP4FakeQuantFunction wraps it. When autograd is on and every expert's weight quantizer is a calibrated dynamic-NVFP4 fake quantizer (block 16, per-tensor amax, no pre-quant scale, rotation, backend or hooks, and for a clipped backward an amax in the weight's dtype), _QuantFusedExperts fake-quantizes all experts of a projection once per forward and gives each expert its slice. Anything else keeps the per-expert path. Outputs and gradients are bit-identical to it. Qwen3-30B-A3B MoE layers, 16K tokens per GPU, 8x B200, FSDP2 as in examples/llm_qat, ms per layer ((8-layer step - 4-layer step) / 4): with activation checkpointing: 426 -> 203 without activation checkpointing: 207 -> 78.5 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Shiyang Chen <shiychen@nvidia.com>
…replace_function fix Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Shiyang Chen <shiychen@nvidia.com>
Contributor
|
No actionable comments were generated in the recent review.
 |
| Layer / File(s) | Summary |
|---|---|
Per-slice FP4 quantization modelopt/torch/kernels/quantization/gemm/fp4_kernel_hopper.py, modelopt/torch/quantization/tensor_quant.py, tests/gpu/torch/quantization/test_tensor_quant_cuda.py |
The FP4 quantizer accepts one global_amax value per first-dimension slice. The new stacked fake-quant function uses this path and the existing gradient helper. CUDA tests compare batched results with separate per-slice calls and check invalid scale counts. |
Fused expert stacked quantization modelopt/torch/quantization/plugins/huggingface.py, tests/gpu/torch/quantization/plugins/test_fused_experts_cuda.py, CHANGELOG.rst |
Fused expert forwards attempt stacked quantization when the quantizers meet the required conditions and otherwise use per-expert quantization. CUDA tests compare loss and gradients against the per-expert path. The changelog describes the feature and reports a speedup measurement. |
Checkpoint recomputation cleanup
| Layer / File(s) | Summary |
|---|---|
Restore patched functions after exceptions modelopt/torch/quantization/utils/core_utils.py, tests/unit/torch/quantization/plugins/test_fused_experts.py, CHANGELOG.rst |
replace_function restores the patched function and removes its cached original in a finally block. A regression test checks that non-reentrant checkpoint recomputation leaves F.linear restored. The changelog records the mixed Tensor/DTensor error. |
Priority:
Estimated code review effort: 3 (Moderate) | ~25 minutes
Change: Feature · Severity of issue fixed: Medium
Sequence Diagram(s)
sequenceDiagram
participant FusedExperts as _QuantFusedExperts
participant QuantizedWeight as _quantized_weight
participant AmaxCheck as _stacked_nvfp4_amax
participant StackedQuant as stacked_nvfp4_fake_quant
participant FP4Kernel as fp4_fake_quant_block
FusedExperts->>QuantizedWeight: request expert projection weight
QuantizedWeight->>AmaxCheck: validate stacked quantizer conditions
AmaxCheck-->>QuantizedWeight: per-expert amax values or None
QuantizedWeight->>StackedQuant: quantize stacked projection weights
StackedQuant->>FP4Kernel: pass stacked weights and per-slice amax
FP4Kernel-->>StackedQuant: return quantized stacked weights
StackedQuant-->>QuantizedWeight: return quantized stacked weights
QuantizedWeight-->>FusedExperts: return selected expert weight
Suggested reviewers: meenchen, chadvoegele
Merge Risk: 
0c247
-
No concrete issue remains that should block merging after normal checks.
Pre-merge checks: 
-
Check name Status Explanation Resolution Docstring Coverage Docstring coverage is 73.91% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 23 functions across 7 files. (1 skipped: … Write docstrings for the functions missing them to satisfy the coverage threshold. -
Check name Status Explanation Description Check Check skipped - CodeRabbit’s high-level summary is enabled. Title check The title clearly and concisely describes the primary change: quantizing all Hugging Face fused NVFP4 experts in one kernel launch. Linked Issues check Check skipped because no linked issues were found for this pull request. Out of Scope Changes check Check skipped because no linked issues were found for this pull request. Security Anti-Patterns PASS. The PR changes four modeloptPython files and noexamplesPython files. Added-code inspection found notorch.load(..., weights_only=False),numpy.load(..., allow_pickle=True), hardcoded …Full details: Security Anti-Patterns
-
Explanation
PASS. The PR changes four
modeloptPython files and noexamplesPython files. Added-code inspection found notorch.load(..., weights_only=False),numpy.load(..., allow_pickle=True), hardcodedtrust_remote_code=True, external-inputeval()/exec(), or# nosec. Nopyproject.tomlor requirements file changed, so no new dependency licensing issue is introduced.
-
Full details: Docstring Coverage
-
Explanation
Docstring coverage is 73.91% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 23 functions across 7 files. (1 skipped: 1 unsupported.)
Resolution
Write docstrings for the functions missing them to satisfy the coverage threshold.
-
Finishing Touches
-
🧪 Generate unit tests (beta)
-
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #3011 +/- ##
==========================================
- Coverage 79.68% 77.46% -2.22%
==========================================
Files 659 661 +2
Lines 74466 78118 +3652
==========================================
+ Hits 59337 60517 +1180
- Misses 15129 17601 +2472
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:
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
This branch has not been deployed
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
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 (QAT speed-up), bug fix
Speed-up. On the eager fused-experts path (
_QuantFusedExperts; ModelOpt forcesexperts_implementation="eager"), each expert's weight goes through its ownTensorQuantizer. That is one Triton launch per expert and projection (256 per layer forward on Qwen3-30B-A3B), plus aweight[idx]slice whose backward writes a gradient the size of all experts once per expert.fp4_fake_quant_blocknow takes one global amax per index of dim 0. One launch quantizes a stacked(E, N, K)weight exactly asEseparate calls would.StackedNVFP4FakeQuantFunctionwraps it with the same backward asdynamic_block_quant(STE, or clipped STE)._QuantFusedExpertsfake-quantizes all experts of a projection on its first call in each forward and gives each expert its slice.unbindwrites one gradient for all experts.It does this only when the result is bit-identical to the per-expert quantizers:
Anything else keeps the per-expert path. That includes no-grad forwards, which quantize only the routed experts, and
FP8Experts.Bug fix.
replace_functionrestored the patched function only when its body exited normally. Non-reentrant activation checkpointing ends its recompute early by raising inside the module forward. So the experts'F.linearinterception stayed installed for the whole process, and the next layer's recompute ran through it. With theexamples/llm_qatFSDP2 config (fsdp_activation_checkpointing: true) on Qwen3-MoE, training crashes withaten.mm.default got mixed torch.Tensor and DTensor. The original is now restored in afinally. #2745 has the same fix (67d8ea0), so whichever PR lands second drops it.Usage
No new API. NVFP4 QAT/QAD of a Hugging Face fused-MoE model uses it automatically:
Testing
ms per MoE layer: Qwen3-30B-A3B, 16K tokens per GPU, 8x B200, FSDP2 as in
examples/llm_qat. Per layer = (8-layer step − 4-layer step) / 4.The per-expert baseline with checkpointing includes the
replace_functionfix, since it crashes without it. Peak memory for 8 layers drops from 13.1 to 11.5 GiB with checkpointing and from 35.5 to 34.5 GiB without.New tests:
tests/unit/torch/quantization/plugins/test_fused_experts.py::test_checkpoint_recompute_restores_linear: runs on CPU in 0.3 s. Fails without thefinally.tests/gpu/torch/quantization/test_tensor_quant_cuda.py::test_fp4_kernel_per_slice_global_amax: a per-slice amax, including 0 and inf, with tiles that span slices, matches per-slice calls in fp32, fp16 and bf16. An amax of the wrong size raises.tests/gpu/torch/quantization/plugins/test_fused_experts_cuda.py: trains a tiny Qwen3-MoE for both backward modes, with and without non-reentrant checkpointing. The loss and every gradient are bit-equal to the per-expert path, and the stacked kernel runs once per projection per forward. The checkpointing case also fails without thefinally.What I ran on B200 (NeMo 26.08 container, transformers 5.18):
test_tensor_quant_cuda.py,test_fp8_experts.pyandtests/unit/torch/quantization/plugins: 374 passed, 2 skipped.Not run: multi-node, an end-to-end
examples/llm_qattraining and accuracy run, and Hopper GPUs.Before your PR is "Ready for review"
Make sure you read and follow Contributor guidelines and your commits are signed (
git commit -s -S).Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded
trust_remote_code=True,torch.load(..., weights_only=False),pickle, etc.).CONTRIBUTING.md: N/AAdditional Information
This PR does not depend on #2745 (the HF EP stack). Its kernel and autograd pieces apply unchanged there.
Once #2745 lands, default QAT runs experts with
grouped_mmand no longer reaches this eager hook. The same one-kernel weight quantization on that path measured 20.1 ms per layer with checkpointing on the same node; Megatron ModelOpt QAD measured 24.6 ms. Adding it is a small follow-up.This PR and #2745 edit the same functions in
huggingface.py, so whichever lands second needs a rebase.🤖 Generated with Claude Code
Summary by CodeRabbit