Repository navigation
Inference: Add support for expert tensor parallelism - #4743
sidsingh-nvidia merged 15 commits into
Conversation
Signed-off-by: Siddharth Singh <sidsingh@nvidia.com>
|
This PR has been automatically converted to draft because all PRs must start as drafts. When you are ready for review, click Ready for Review to begin the review process. This will:
See the contribution guide for more details. |
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
Signed-off-by: Siddharth Singh <sidsingh@nvidia.com>
|
Hi, I’d like to continue with multi-GPU L20 validation of the NCCL path for ETP. |
|
Status update from the L20 validation side, so this does not stay silent. I checked the current base ( # megatron/core/transformer/transformer_config.py:1856
if self.expert_tensor_parallel_size > 1:
raise ValueError(
"Inference-optimized MoE layers does not support expert tensor parallelism."
)and that this PR's Rather than duplicate this PR, I opened #7484 for the design question (which axis the fused One finding that may be useful here: If you would rather this work continue on top of your branch, or if the sharding axis is |
…parallel Merge upstream main (f5d4058) into the ETP inference branch. Conflict resolution: - dynamic_context.py: keep main's dispatcher buffer allocation (unconditional valid-tokens tensor, nccl/nvls branches, vLLM fused-MoE buffers); size the NVLS buffers and the new _nvls_dispatcher flag on the tp_ep group, and size the vLLM MoE buffers on tp_ep rows with fc1 width divided by ETP. - dynamic_engine.py: keep the tp_ep consensus naming; take main's local_schedulable work count; fix main's new ep_world_size reference. - token_dispatcher_inference.py: take main's mask_routing_padding / ep_size==1 early-out and metadata code; apply tp_size * ep_size scaling. - transformer_config.py: take main's backend/mxfp8 setup; keep the branch's removal of the ETP>1 rejection. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Siddharth Singh <sidsingh@nvidia.com>
Follow-up to the main merge, reducing the branch to a minimal NVLS-only ETP > 1 inference path: - transformer_config: replace the commented-out ETP check with a fail-fast ValueError when expert_tensor_parallel_size > 1 is used with a non-NVLS inference MoE dispatcher (NCCLAllGatherDispatcher does not support ETP). - dynamic_context: add the tp_ep group as a separate attribute (pg_collection, with an MPU migration fallback that also works for EP == 1); use it only for the NVLS dispatcher flag, NVLS buffer sizing and vLLM MoE buffer sizing. The nccl/training dispatcher flags stay on the EP group as on main. Drop the set_tp_ep_zmq_communicator rename. - dynamic_engine: revert the tp_ep renames (tests rely on the main names) and instead run the existing EP consensus/ZMQ communicator over the tp_ep group. - token_dispatcher_inference: formatting/comments only. - attention.py: drop the debug `HAVE_FA4 = False` override. - tests: update NVLSAllGatherVDispatcher.allocate_buffers kwarg to tp_ep_group. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Siddharth Singh <sidsingh@nvidia.com>
…p_world_size These attributes now hold the rank and size of the expert tensor-and-model-parallel (tp_ep) group, not the EP group, so name them after it. They are read only inside DynamicInferenceEngine (ZMQ consensus communicator setup and the consensus all-reduce) and by the coordinator test's DummyEngine. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Siddharth Singh <sidsingh@aws-cmh-slurm-1-vscode-02.cm.cluster>
…ilogue The vLLM-backend fused MoE now multiplies each row by its routing probability inside the FC2 grouped GEMM (MUL_ROUTED_WEIGHT, on the fp32 accumulator before the bf16 store), as upstream vLLM does, instead of in _moe_sum. _moe_sum is now an unweighted top-k sum. Batch-invariant mode is unchanged (probabilities are applied at the activation there). GSM8K (1319, greedy) vs main EP4/ETP1: EP4/ETP1 87.6%, EP2/ETP2 87.7%, EP1/ETP4 87.2% (main 87.1%). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Siddharth Singh <sidsingh@aws-cmh-slurm-1-vscode-02.cm.cluster>
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/40081499889 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/40086697791 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/40235043205 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/40319495733 |
Signed-off-by: Siddharth Singh <sidsingh@nvidia.com>
|
/ok to test b6ec6e2 |
|
/ok to test 9747936 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/40539511053 |
`gpt_dynamic_inference_tp4_pp1_ep4_16B_inference_optimized_swiglu_logitsmatch`: NVIDIA#4743 changed this test's logprobs on both platforms and regenerated the GB200 goldens on main after the fp32-logits refresh had rewritten them, so the two could not be merged; the H100 goldens were left stale on main. Both regenerated from a run of this branch on main `2abdd0d1c6` (GitLab pipeline 77095337, jobs 512966250 and 512966248). Generated tokens unchanged; GB200 logprobs within 0.05 of main's new values. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Main NVIDIA#4743 makes every context whose expert tensor-and-model-parallel group spans ranks use the NVLS dispatcher, which wires the context's real-token count into a class attribute. The EP2 contexts built in test_dynamic_context.py left it set, so the later NVLS tests masked their routing rows against a stale count and the graph replay differed. Clear the dispatcher's class state before the NVLS tests set up. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Lawrence McAfee <lmcafee@nvidia.com>
What does this PR do ?
Adding support for expert tensor parallelism in inference.
Todos:
Issue tracking
For PRs from open-source community contributors:
Linked issue:
Contribution process
Pre-checks
Code review
Feel free to message or comment the @mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!
All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.
Step 1: Mark PR as "Ready for Review"
.github/CODEOWNERS.Final Review might get declined if these requirements are not fulfilled.
Step 2: Final Review
For PRs that change
megatron/core, once all expert reviewers have approved, theFinal Reviewlabel is applied automatically and final reviewers are assigned.For PRs outside
megatron/core, this step is skipped.Step 3: Approved
Once all required reviewers have approved, the
Approvedlabel is applied automatically.Merge
Any member of mcore-engineers will be able to merge your PR.
For MRs into `dev` branch
The proposed review process for `dev` branch is under active discussion.MRs are mergable after one approval by either
eharper@nvidia.comorzijiey@nvidia.com.