fix: EP collective deadlock with variable-length token counts (LoRA flavor) - #3631
Open
akx wants to merge 1 commit into
Open
fix: EP collective deadlock with variable-length token counts (LoRA flavor)#3631akx wants to merge 1 commit into
akx wants to merge 1 commit into
Conversation
…lavor) GroupedExpertsLoRA.forward still gathers tokens across the EP mesh with DTensor.from_local(x, [Shard(0)]).full_tensor(), which assumes uniform token counts across ranks. PR NVIDIA-NeMo#1365 (8f2b685) fixed this in GroupedExperts but not in the LoRA subclass, so LoRA on MoE experts with unpacked or unpadded batches deadlocks in NCCL on the first MoE layer. Port of NVIDIA-NeMo#1365's fix to lora_experts.py: pad+all_gather+trim on the way in, all_reduce+narrow on the way out, gradient anchor so every rank enters the backward collectives. Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Aarni Koskela <akx@iki.fi>
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 does this PR do?
Fixes NCCL communication deadlocks when using Expert Parallelism when training LoRAs.
Port of #1365's fix (8f2b685) to lora_experts.py, and this subclass.
Changelog
See #1365; this is a port of it.
Before your PR is "Ready for review"
Pre checks:
If you haven't finished some of the above items you can still open "Draft" PR.
Additional Information
See the below block for Claude Fable's analysis.
Details
**Symptom:** Expert-parallel LoRA training with unpacked, unpadded batches (`packed_sequence_size: 0`, `padding: false`, `local_batch_size: 1`) hangs forever on the first MoE layer of the first forward. Every rank sits in the same `all_gather_into_tensor` (native stack: `ProcessGroupNCCL::allgather_into_tensor_coalesced` issued from `lora_experts.py:184` via `DTensor.full_tensor()`), ~200 W per GPU, no error, no timeout until the NCCL heartbeat. Reproduced on 1×8 and 2×8 nodes, `experts: torch` and `torch_mm`, with and without activation checkpointing / dynamo.Root cause: the LoRA forward gathers tokens across the EP mesh with
from_local(..., Shard(0))does not exchange sizes — it assumes every rank's local tensor has the same shape and infers the global shape aslocal × ep_size. With ragged per-rank token counts (e.g. 68 / 401 / 1607 tokens across ranks), each rank issues a differently sized all-gather, which NCCL cannot detect; it waits forever. The non-LoRA parentGroupedExperts.forwardhandles this correctly (exchanges lengths, pads,_AllGatherConcatVarlenFn, then all-reduce + narrow to combine). The subclass docstring says it "mirrors GroupedExperts.forward" but uses the naive DTensor gather. It only works when batches are packed or padded — which every shipped recipe does, so the bug is invisible in-tree.Repro: any MoE LoRA config with
ep_size > 1,packed_sequence_size: 0,padding: false,local_batch_size: 1, on a dataset with variable-length samples (we usedHuggingFaceH4/no_robotswith Kimi-K2.5; a 2-layer slice of the checkpoint reproduces it in under 2 minutes on one node). Log ax.shape[0]per rank before the gather to see the mismatch.