[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, which sizes dispatch/combine output buffers dynamically to the per-step recv-token total (instead of the static
Confidence Score: 4/5The new eager/overflow modes are internally consistent and well-guarded, but the public API changes — positional parameter reorder in ep_bootstrap/EpBuffer.init and the token_counts→tokens_per_expert rename with dtype change — silently break positional callers and attribute readers, despite being marked non-breaking. The core eager/overflow logic is correct: mutual-exclusivity enforced at both layers, host-sync for dynamic sizing is intentional and documented, the backward is correctly reshaped. The unresolved breaking changes to positional argument order and the public EpBuffer attribute are the primary risk for downstream integrations. Files Needing Attention: transformer_engine/pytorch/ep.py — public API surface changed in ways that silently break positional callers and attribute readers Important Files Changed
Sequence DiagramsequenceDiagram
participant Caller
participant ep_dispatch
participant ep_prepare
participant NCCL_EP
participant _EpDispatch
Note over Caller,_EpDispatch: Eager mode (recv_capacity_per_rank=None)
Caller->>ep_dispatch: ep_dispatch(buffer, tokens, topk_idx, topk_w)
ep_dispatch->>ep_prepare: ep_prepare(buffer, topk_idx)
ep_prepare->>NCCL_EP: ncclEpUpdateHandle (AllGather routing)
NCCL_EP-->>ep_prepare: total_recv_tokens [device int64]
ep_prepare->>ep_prepare: "buffer._host_total_recv_tokens = .item() [D2H sync]"
ep_prepare-->>ep_dispatch: tokens_per_expert (buffer tensor)
ep_dispatch->>ep_dispatch: "rows = buffer._host_total_recv_tokens"
ep_dispatch->>ep_dispatch: alloc recv_tokens[rows,H] and recv_w[rows]
ep_dispatch->>_EpDispatch: apply(handle_mem, recv_tokens, recv_w, topk_idx, tokens, topk_w)
_EpDispatch->>NCCL_EP: ncclEpDispatch
NCCL_EP-->>_EpDispatch: recv buffers filled
_EpDispatch-->>ep_dispatch: (recv_tokens, recv_w)
ep_dispatch-->>Caller: (recv_tokens, recv_w, tokens_per_expert)
Note over Caller,_EpDispatch: Non-eager mode (fixed recv_capacity_per_rank)
Caller->>ep_dispatch: ep_dispatch(buffer, tokens, topk_idx, topk_w)
ep_dispatch->>ep_prepare: ep_prepare(buffer, topk_idx)
ep_prepare->>NCCL_EP: ncclEpUpdateHandle
NCCL_EP-->>ep_prepare: tokens_per_expert and total_recv_tokens [device only]
ep_prepare-->>ep_dispatch: tokens_per_expert
ep_dispatch->>ep_dispatch: "rows = buffer.recv_capacity_per_rank (static)"
ep_dispatch->>_EpDispatch: apply(...)
_EpDispatch-->>ep_dispatch: (recv_tokens, recv_w)
ep_dispatch-->>Caller: (recv_tokens, recv_w, tokens_per_expert)
Reviews (4): Last reviewed commit: "Merge branch 'main' into pyt_ep_eager" | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| @@ -95,6 +98,9 @@ def ep_bootstrap( | |||
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.
There was a problem hiding this comment.
We can do that too. Let me update the design.
| 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.
There was a problem hiding this comment.
NCCL EP could work with the scenario in which there are multiple num_topk (in different layers), which we may not need.
We can simplify this to num_topk and make it mandatory.
| 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?
There was a problem hiding this comment.
For drop_on_overflow, we need recv_capacity which eager mode does not provide.
…nk, require num_topk, and run prepare in ep_dispatch Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
…ecv-buffer API Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
|
/te-ci L1 |
| alignment = int(alignment) | ||
| if alignment > 1 and (alignment & (alignment - 1)) != 0: | ||
| raise ValueError(f"alignment must be 0, 1, or a power of two (got {alignment}).") | ||
| self.eager = _EAGER |
There was a problem hiding this comment.
This adds a dependency that buffer creation needs to happen after bootstrap, can we add assertion for it?
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
|
/te-ci L1 |
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: