fix(mmlu): shard batches over expert-DP group for MoE models#1966
Closed
ChenhanYu wants to merge 1 commit into
Closed
fix(mmlu): shard batches over expert-DP group for MoE models#1966ChenhanYu wants to merge 1 commit into
ChenhanYu wants to merge 1 commit into
Conversation
megatron_mmlu shards whole batches across the dense data-parallel group and has each rank run megatron_prefill on a disjoint subset. That is correct for dense models, but for MoE models with expert parallelism (EP>1) the prefill forward runs an expert all-to-all across the EP group. When EP overlaps the dense-DP group, the per-rank disjoint batches desync that all-to-all — uneven batch counts and differing padded seq-lengths leave trailing ranks blocked at NCCL communicator creation until the store timeout. Shard over the expert-data-parallel group instead when EP>1, so ranks within an EP group evaluate every batch in lockstep and only true expert-DP replicas take disjoint batches. Dense models (EP==1) keep the standard DP sharding. 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)
📝 WalkthroughWalkthroughThe MMLU plugin now selects expert-data-parallel sharding state when expert model parallelism is active and retains standard data-parallel sharding otherwise. ChangesMMLU sharding
Estimated code review effort: 2 (Simple) | ~10 minutes 🚥 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}
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1966 +/- ##
==========================================
- Coverage 77.92% 77.29% -0.64%
==========================================
Files 519 519
Lines 58217 58221 +4
==========================================
- Hits 45366 45002 -364
- Misses 12851 13219 +368
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). |
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.
Problem
megatron_mmlushards whole batches across the dense data-parallel group (get_data_parallel_world_size()), giving each rank a disjoint subset (if batch_idx % dp_size != dp_rank: continue) and all-reducing per-subject counts over the DP group at the end.That is correct for dense models — DP ranks are independent replicas and the forward has no cross-rank collective. It deadlocks for MoE models with expert parallelism (EP>1):
megatron_prefill's 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=1on 4 GPUs — all four ranks are both EP peers and dense-DP ranks), each rank runs prefill on a different batch, so their expert all-to-alls desync. Uneven batch counts (n_batches % dp_size != 0) and per-batch padded seq-lengths leave trailing ranks blocked at NCCL communicator creation until the c10d store times out (600 s).Reproduced with a NemotronH-class MoE (Nano-30B-A3B) quantize → MMLU at
EP=4: ranks 2 & 3 wait on rank 0'sncclUniqueIdand time out.Fix
Shard over the expert-data-parallel group when
EP>1. Ranks within one EP group then evaluate every batch in lockstep (so the expert all-to-all always has all participants), and only true expert-DP replicas — each holding a full expert set — take disjoint batches. Dense models (EP==1) keep the standard DP sharding unchanged.EP=4on 4 GPUs → expert-DP=1 → no sharding, all ranks lockstep. Final all-reduce becomes a size-1 no-op (correct).EP=4on 8 GPUs → expert-DP=2 → two replicas shard batches, each replica's 4 EP ranks lockstep.Test plan
EP>1) completes instead of hanging — validated on Nano-30B-A3BEP=4.EP==1path is identical to before).🤖 Generated with Claude Code
Summary by CodeRabbit