example: Nemotron-3-Nano-30B-A3B-BF16 Megatron-LM PTQ launcher example#1964
Closed
ChenhanYu wants to merge 1 commit into
Closed
example: Nemotron-3-Nano-30B-A3B-BF16 Megatron-LM PTQ launcher example#1964ChenhanYu wants to merge 1 commit into
ChenhanYu wants to merge 1 commit into
Conversation
Adds a quantize + MMLU-gate (EP=4) -> export (PP=4) -> vLLM smoke pipeline for the smallest NemotronH-class MoE, modeled on the Super-120B-A12B megatron_lm_ptq example. The quantize and MMLU stages run on the MoE expert-parallel (EP=4) collective; export uses PP=4 for the hybrid layer stack. Uses NVFP4_DEFAULT_CFG (no Nano-30B recipe exists) and sets modelopt_install_path to the nemo-container venv so the mounted modelopt overrides the container's copy. Signed-off-by: Chenhan Yu <chenhany@nvidia.com>
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 (1)
📝 WalkthroughWalkthroughAdds a Slurm configuration for NVIDIA Nemotron-3-Nano PTQ using NVFP4, including quantization with an MMLU gate, pipeline-parallel checkpoint export, and vLLM smoke testing of the exported checkpoint. ChangesNemotron-3 Nano PTQ workflow
Estimated code review effort: 2 (Simple) | ~10 minutes Sequence Diagram(s)sequenceDiagram
participant Slurm
participant quantize.sh
participant export.sh
participant vLLM
Slurm->>quantize.sh: Quantize checkpoint and run MMLU gate
quantize.sh-->>export.sh: Provide quantized checkpoint
export.sh->>export.sh: Export NVFP4 checkpoint with PP=4
export.sh-->>vLLM: Provide exported checkpoint
vLLM->>vLLM: Run GPQA sample queries with TP=4
🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #1964 +/- ##
==========================================
+ Coverage 77.92% 77.95% +0.03%
==========================================
Files 519 519
Lines 58217 58217
==========================================
+ Hits 45366 45385 +19
+ Misses 12851 12832 -19
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:
|
3 tasks
Collaborator
Author
|
Superseded by #1967, which combines the launcher example with the MMLU expert-DP sharding fix (the example is the fix's end-to-end test, so they ship together). |
Contributor
|
ChenhanYu
added a commit
that referenced
this pull request
Jul 14, 2026
#1967) Combines the NemotronH MoE PTQ launcher example with the MMLU expert-parallel sharding fix it exercises. The example is the end-to-end test for the fix — quantize + **MMLU EP=4** + export + vLLM smoke on Nano-30B-A3B — so they ship together. ## Fix — `megatron_mmlu.py` shards over the expert-DP group for MoE models `megatron_mmlu` shards whole batches across the **dense data-parallel group** and runs `megatron_prefill` per-rank on a disjoint subset. Correct for dense models, but for MoE with **EP>1** the prefill forward runs an **expert all-to-all across the EP group**. When EP overlaps the dense-DP group (e.g. `EP=4,TP=1,PP=1` on 4 GPUs), each rank is on a different batch, so the all-to-alls desync (uneven batch counts + differing padded seq-lengths) → trailing ranks block at NCCL communicator creation until the c10d store times out (600 s). Reproduced on Nano-30B-A3B at `EP=4`: ranks 2 & 3 wait on rank 0's `ncclUniqueId`. The fix shards over the **expert-data-parallel** group when `EP>1`, so EP peers evaluate every batch in lockstep and only true expert-DP replicas take disjoint batches. Dense models (`EP==1`) are byte-for-byte unchanged. ## Example — `examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/megatron_lm_ptq.yaml` Quantize (EP=4) + MMLU gate (EP=4) → export (PP=4) → vLLM smoke, modeled on the Super-120B example. Uses `NVFP4_DEFAULT_CFG` (no Nano-30B recipe) and sets `modelopt_install_path` to the nemo-container venv so the mounted modelopt overrides the container copy. Passes `test_examples_resolve.py`. ## Test plan - [x] MoE MMLU (`EP>1`) completes instead of hanging — validated end-to-end via this example on Nano-30B-A3B `EP=4` (nmm-sandbox CI). - [x] `test_examples_resolve.py` green (example parses/resolves). - [ ] Dense MMLU sharding unchanged (`EP==1` path identical). Supersedes #1964 (example) and #1966 (fix). 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Improved MMLU evaluation for expert-parallel MoE models by aligning batch sharding with expert-parallel execution for more consistent results. * **Bug Fixes** * Dense (non-MoE) models retain standard data-parallel batch sharding, while MoE handling is corrected. * **Documentation** * Updated the Nemotron-3-Nano BF16 PTQ example to clarify the quantize → MMLU gate → export flow and why the vLLM smoke is intentionally omitted. * Adjusted the Llama-3.2-1B-Instruct PTQ/MMLU gate lower-bound threshold (0.40 → 0.36). * **Tests** * Marked the homogeneous compressed sharded state-dict test to skip on Blackwell to avoid a known flaky issue. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Chenhan Yu <chenhany@nvidia.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.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
1 participant
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
Adds a Megatron-LM PTQ launcher example for NVIDIA-Nemotron-3-Nano-30B-A3B-BF16, the smallest NemotronH-class hybrid Mamba-Transformer MoE we ship an example for. Modeled on the existing
NVIDIA-Nemotron-3-Super-120B-A12B-BF16/megatron_lm_ptq.yaml.Pipeline (1 node × 4 GPUs throughout):
Notes
NVFP4_DEFAULT_CFG— there is no dedicated Nano-30B recipe (only Nano-4B / Super-120B / Ultra-550B).modelopt_install_pathto the nemo-container venv (/opt/venv/.../site-packages/modelopt) so the mounted modelopt overrides the container's copy — the container's active interpreter is a venv whose site-packages precede/usr/localonsys.path, so the defaultdist-packagesmount is otherwise ignored. Matches the Super-120B example's convention.tools/launcher/tests/test_examples_resolve.py(CPU-only parse/resolve gate).Test plan
test_examples_resolve.pygreen (CI)🤖 Generated with Claude Code
Summary by CodeRabbit