[PyTorch] NCCL EP eager mode and drop-on-overflow policy#3229
Conversation
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
…cv_tokens Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
Greptile SummaryThis PR adds two new NCCL EP dispatch modes to the PyTorch API: eager mode (sizes recv buffers per-step from the actual routing recv-token total, requiring a D2H sync) and drop-on-overflow (silently drops tokens that exceed
Confidence Score: 3/5Merging introduces a breaking rename of a public EpBuffer attribute and dtype change that will silently fail for any downstream user before review/migration guidance is added. The C++ and binding layers are cleanly implemented with proper mutual-exclusion guards. The Python eager-mode flow is logically sound. However, renaming EpBuffer.token_counts to tokens_per_expert and changing its dtype from int32 to int64 breaks existing callers — the PR declares itself non-breaking, so this may have been an oversight. transformer_engine/pytorch/ep.py — the EpBuffer attribute rename/dtype change and the _EpDispatch.forward dead field need attention before merge. Important Files Changed
Sequence DiagramsequenceDiagram
participant User
participant ep_dispatch
participant ep_prepare
participant NCCL_EP as NCCL EP (device)
participant Host
Note over User,Host: Eager mode (buffer.eager=True)
User->>ep_dispatch: ep_dispatch(buffer, tokens, topk_idx, topk_weights)
ep_dispatch->>ep_prepare: ep_prepare(buffer, topk_idx)
ep_prepare->>NCCL_EP: ncclEpUpdateHandle (AllGather routing)
ep_prepare->>Host: total_recv_tokens.item() [D2H sync]
Host-->>ep_prepare: rows (int)
ep_prepare-->>ep_dispatch: "buffer._host_total_recv_tokens = rows"
ep_dispatch->>ep_dispatch: alloc recv_tokens(rows, H), recv_topk_weights(rows,)
ep_dispatch->>NCCL_EP: "ncclEpDispatch (skip_prepare=True)"
ep_dispatch-->>User: (recv_tokens[rows,H], recv_topk_weights[rows], tokens_per_expert)
Note over User,Host: Non-eager / drop-on-overflow mode
User->>ep_dispatch: ep_dispatch(buffer, tokens, topk_idx, topk_weights)
ep_dispatch->>ep_dispatch: "rows = buffer.recv_capacity_per_rank"
ep_dispatch->>NCCL_EP: ncclEpUpdateHandle (AllGather routing)
Note right of NCCL_EP: drop-on-overflow: excess tokens dropped silently
ep_dispatch->>NCCL_EP: ncclEpDispatch
ep_dispatch-->>User: (recv_tokens[capacity,H], recv_topk_weights[capacity], tokens_per_expert)
|
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| @@ -409,13 +472,13 @@ | |||
| ctx.top_k = topk_weights.shape[-1] | |||
| ctx.recv_capacity = recv_tokens.shape[0] | |||
There was a problem hiding this comment.
ctx.recv_capacity is saved but never read in backward
ctx.recv_capacity = recv_tokens.shape[0] stores the (possibly dynamic in eager mode) number of recv slots, but nothing in backward reads ctx.recv_capacity. The backward reconstructs grad shapes entirely from ctx.tokens_T_flat, ctx.hidden_dim, etc. This attribute is dead code and could mislead a future reader into thinking it is consumed by the backward kernel.
| if buffer.eager: | ||
| if recv_tokens is not None or recv_topk_weights is not None: | ||
| raise ValueError( | ||
| "eager mode sizes the recv outputs from the per-step recv-token total " | ||
| "and cannot use caller-supplied recv_tokens / recv_topk_weights" | ||
| ) | ||
| # Prepare first to learn this step's recv total, then size the recv | ||
| # outputs to it; skip re-running prepare inside the autograd forward. | ||
| ep_prepare(buffer, topk_idx) | ||
| rows = buffer._host_total_recv_tokens | ||
| skip_prepare = True |
There was a problem hiding this comment.
ep_prepare + ep_dispatch in eager mode
In eager mode, ep_dispatch always calls ep_prepare internally (line 627) to learn the per-step recv total before sizing the output buffers. If a user has already called ep_prepare(buffer, topk_idx) externally (e.g., to inspect tokens_per_expert before dispatching), prepare runs twice on the same handle_mem in the same step. The second call is idempotent in normal operation, but it duplicates the NCCL AllGather and the D2H sync, wasting latency. The docstring for ep_dispatch does not warn against this. A short note — "In eager mode ep_dispatch runs ep_prepare internally; do not call ep_prepare separately before ep_dispatch" — would prevent the confusion.
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
| ep_group: dist.ProcessGroup, | ||
| num_experts: int, | ||
| max_tokens_per_rank: int, | ||
| recv_capacity_per_rank: int, |
There was a problem hiding this comment.
With eager, do we still need to have recv_capacity_per_rank as a always required arg? I'm thinking to make it a kwarg, if user feed it, we mark eager to False, otherwise mark True. And we can get rid of eager kwarg.
| bound. This requires a host sync each step, so it is not CUDA-graph | ||
| capturable. Mutually exclusive with ``zero_copy``. Defaults to ``False``. | ||
|
|
||
| ``max_num_topk`` is the upper bound on per-token top-k; it sizes NCCL EP |
There was a problem hiding this comment.
Not sure I understand this, is just just the top_k? upper bound on per-token top-k is a bit confusing. If it is just topk, I would recommend to make it a arg as it is something constant that user can always provide. User does not need to be aware of the logic that only eager mode requires the topk.
| topk_idx: torch.Tensor, | ||
| topk_weights: torch.Tensor, | ||
| *, | ||
| recv_tokens: Optional[torch.Tensor] = None, |
There was a problem hiding this comment.
Should we assert recv_tokens/recv_topk_weights None in eager mode? I see we have relevant assert in combine
| ) | ||
| return buffer.token_counts | ||
| if buffer.eager: | ||
| buffer._host_total_recv_tokens = int(buffer.total_recv_tokens.item()) |
There was a problem hiding this comment.
Why do we need this buffer._host_total_recv_tokens? Can't we just call item whenever we need?
| f"topk_weights must be float32; got dtype={topk_weights.dtype}. " | ||
| "Cast with topk_weights.float() before calling." | ||
| ) | ||
| skip_prepare = False |
There was a problem hiding this comment.
This skip_prepare logic looks a bit awkward. Can we just move the ep_prepare from the autograd function to here always? Seems to me it is ok, in this way we also do not need to pass buffer.total_recv_tokens to autograd fn
| raise ValueError("ep_bootstrap: zero_copy and eager modes are mutually exclusive") | ||
| if eager and max_num_topk < 1: | ||
| raise ValueError("ep_bootstrap: eager mode requires max_num_topk >= 1") | ||
| if eager and drop_on_overflow: |
There was a problem hiding this comment.
I wonder if we really need to assert this? Overflow should be irrelevant to eager right? When it is in eager mode, drop_on_overflow should be ignored automatically?
2 participants
Description
Add two NCCL EP dispatch options to the PyTorch API:
recv_capacity_per_rank.max_recv_tokens_per_rankare dropped and dispatch continues, instead of trapping.Type of change
Changes
ep.h/ep_backend.cpp: adddrop_on_overflowgroup policy and wirenum_topk+ AUTO recv budget through the group config.pytorch/ep.py:ep_bootstrap(eager=...)toggle captured at bootstrap;EpBufferreads it to size outputs from the per-step recv-token total and exposesnum_recv_tokens. Eager and zero-copy are mutually exclusive, and eager rejects caller-provided output buffers.pytorch/csrc/extensions: plumb the new fields through the dispatch binding.Checklist: