Add an opt-in DeepEP transport for the AutoEP expert all-to-all - #8213
Add an opt-in DeepEP transport for the AutoEP expert all-to-all#8213yh0903 wants to merge 3 commits into
Conversation
The expert all-to-all is the largest single cost in an AutoEP step. Replaying
real SFT routing on 16 H100s across two nodes, NCCL spends 99.7 ms per step on
payload all-to-all against DeepEP's 48.0 ms, and the gap widens with routing
skew: at the most imbalanced step recorded, NCCL degrades to 115.8 ms while
DeepEP stays at 46.4 ms.
Roughly half of that is deduplication. DeepEP sends a token once per
destination rank rather than once per selected expert, worth about 1.29x on
its own; its kernels account for the remaining 1.61x. The skew immunity comes
entirely from the kernels.
In whole training steps the step time drops from 454.0 ms to 370.7 ms, an
18.3% reduction, with the payload all-to-all falling from 31.5% of the step to
8.0%.
DeepEP replaces more than the two collectives. It takes tokens before top-k
expansion and replicates them itself, groups arrivals by expert for the
grouped GEMM, and reduces the weighted sum in its combine, so the expansion
and reduction around the collectives are replaced too. Its backward pass has
no separate entry points: the gradient of a combine is a dispatch and the
gradient of a dispatch is a combine, both replayed against the handle the
forward dispatch produced.
The transport is selected by environment variable and defaults to NCCL:
DEEPSPEED_AUTOEP_COMM_BACKEND=nccl (default)
DEEPSPEED_AUTOEP_COMM_BACKEND=deepep
DEEPSPEED_AUTOEP_COMM_SMS=<n> (default 12)
A job that sets nothing behaves exactly as before, deep_ep is imported only
when it is selected, and an unparsable value warns and falls back rather than
breaking a path nobody opted into. DeepEP also needs NCCL 2.30.4 or newer for
GIN, which not every cluster has, so a missing package explains what it
requires instead of failing deep inside buffer construction.
The default SM budget of 12 was chosen by sweeping whole steps rather than the
collective alone. Communication competes with the expert GEMM for SMs: at 8
SMs the collective itself degrades, and above 12 the step grows because
communication takes SMs the rest of the step was using. The measured steps
were 340, 311, 353, 360 and 391 ms at 8, 12, 16, 24 and 32 SMs.
Signed-off-by: yh0903 <helloyu0903@gmail.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 8e3325bba1
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| def backward(ctx, grad_received): | ||
| grad_tokens = ctx.exchange.combine(grad_received.contiguous(), ctx.handle) | ||
| return None, _conform_rows(grad_tokens, ctx.tokens_shape), None, None |
There was a problem hiding this comment.
When training with the DeepEP backend, the original router scores enter the custom autograd function as topk_weights, but this backward unconditionally returns None for that input. The received weights later read from exchange.last_recv_weights were produced inside the custom forward and are outside its autograd outputs, while _DeepEPCombine.backward also returns no weight gradient, so the MoE loss silently stops contributing gradients to the router gate for both pre- and post-score modes.
Useful? React with 👍 / 👎.
| if self.comm_backend == DEEPEP_BACKEND: | ||
| expert_output = self._deepep_route(x, ro) |
There was a problem hiding this comment.
When AutoEP is used with folded TP (tp_size > 1) and DeepEP is selected, this call ignores the assignment partition created above and returns a fully combined [T, H] tensor on every TP lane. The subsequent restore_combined expects one row per lane-local assignment, gathers those rows, applies routing weights, and scatters by assignment metadata, so it either encounters incompatible row counts or weights and scatters already-combined token outputs as if they were assignment outputs, corrupting the result.
Useful? React with 👍 / 👎.
| num_experts=self.num_experts, | ||
| top_k=self.top_k, | ||
| hidden_size=self.hidden_size, | ||
| num_max_tokens_per_rank=tokens.shape[0], |
There was a problem hiding this comment.
When a later forward has more tokens than the first forward, the retained exchange is still configured with the first batch's token count as num_max_tokens_per_rank. Because the buffer is only created when _deepep_exchange is None, variable sequence lengths, dynamic batching, or a short warm-up batch followed by a larger batch exceed its declared capacity and cause the DeepEP dispatch to fail; recreate or grow the exchange when tokens.shape[0] exceeds the stored capacity.
Useful? React with 👍 / 👎.
| elif self.comm_backend == DEEPEP_BACKEND: | ||
| # Already reduced over top-k and back in token order. | ||
| output = expert_output.reshape(bsz, seqlen, hdim) |
There was a problem hiding this comment.
When DEEPSPEED_AUTOEP_COMM_BACKEND=deepep is set but ep_size == 1, the DeepEP route is never called and expert_output still contains one row per top-k assignment from the local path. This backend-only branch nevertheless skips combine_from_routed and reshapes those T * top_k rows directly to [B, S, H], which fails for top_k > 1 and omits the routing reduction even when the shape happens to fit.
Useful? React with 👍 / 👎.
| @@ -0,0 +1,293 @@ | |||
| # SPDX-License-Identifier: Apache-2.0 | |||
There was a problem hiding this comment.
This is a non-merge commit, but its commit message has no Signed-off-by trailer, so it does not satisfy the repository's mandatory commit and CI requirement. Recreate the proposed commit with --signoff using the configured Git identity.
AGENTS.md reference: AGENTS.md:L8-L8
Useful? React with 👍 / 👎.
| DEEPSPEED_AUTOEP_COMM_BACKEND=nccl (default) | ||
| DEEPSPEED_AUTOEP_COMM_BACKEND=deepep |
There was a problem hiding this comment.
The new opt-in backend, its environment variables, dependency, NCCL requirement, and supported configurations are only described in an internal module docstring; the commit adds no user-facing documentation, so users cannot reliably discover or configure the feature. The repository explicitly requires documentation updates for new features.
AGENTS.md reference: AGENTS.md:L26-L26
Useful? React with 👍 / 👎.
df6e5d4 to
e566d6f
Compare
Six issues raised in review, all confirmed against the code: Router scores stopped receiving gradients. Both autograd functions returned None for their topk_weights input, and the received weights were read off the exchange rather than returned as an output, leaving them outside the graph. The MoE loss therefore contributed nothing to the router gate in either score mode. Dispatch now returns the received weights so autograd carries their gradient back, and both backward passes reduce and return it. This fails silently -- training runs and the loss falls while the gate never learns -- so two tests pin the return arity that carries it. The buffer kept the first batch's token count as its capacity. Variable sequence lengths or a small warm-up batch followed by a larger one would exceed it; the exchange is now rebuilt when a batch outgrows it. Skipping the routed combine was keyed on the backend being selected rather than on DeepEP having run. With ep_size == 1 the local path runs instead and still has one row per assignment, so the reduction was skipped on rows that needed it. It is now keyed on the route actually producing the output. Folded tensor parallelism partitions assignments across lanes and restores them by assignment metadata, which a transport whose combine returns token-major rows cannot satisfy. Such layers now warn and use the NCCL path rather than corrupting the result. The backend, its environment variables, its NCCL requirement and its limits are now documented in the AutoEP page rather than only in a module docstring. Signed-off-by: yh0903 <helloyu0903@gmail.com>
2 participants
What this adds
An optional DeepEP transport for the AutoEP expert all-to-all, selected by environment variable and defaulting to the existing NCCL path.
Why
The expert all-to-all is the largest single cost in an AutoEP step. Replaying routing captured from real SFT runs on 16 H100s across two nodes:
Two things are worth separating. Roughly half the gain is deduplication: DeepEP sends a token once per destination rank rather than once per selected expert, which is worth about 1.29x on its own and could in principle be done without changing transport. Its kernels account for the remaining 1.61x, and they are also where the skew immunity comes from — at the most imbalanced step recorded, NCCL degrades to 115.8 ms while DeepEP stays flat.
In whole training steps, with the same model, data and step count and only the backend changed:
Why DeepEP v2 only
This wraps
ElasticBuffer, the v2 API. The legacy v1Bufferis deliberately not supported:the
NVreg_EnableStreamMemOPsdriver parameter or the GDRCopy/dev/gdrdrvdevice. Neither is present on the cluster this was developed against, and both require administrator action rather than configuration.against 153 GB/s intranode — and the node boundary is exactly the cost this change exists to reduce. v2's hierarchical NVLink plus RDMA path targets that case directly.
that already has ZeRO and data-parallel groups on the fabric.
The backend name is
deepeprather thandeepep_v2: it names the library, and nothing about it would have to change if v1 were ever added.Choosing the SM budget
Communication competes with the expert GEMM for SMs, so the default was chosen by sweeping whole training steps rather than the collective in isolation:
At 8 the collective itself degrades; above 12 the step grows because communication takes SMs the rest of the step was using. Tuning this from 24 to 12 moved the end-to-end result from 1.126x to 1.225x.
Scope and safety
git diff -wremoves zero lines from the shipped path: the existing NCCL code is unchanged and only moves into anelsebranch.deep_epis imported only when the backend is selected, so installations without it are unaffected, and an unparsable backend value warns and falls back rather than breaking a path nobody opted into. DeepEP also requires NCCL 2.30.4 or newer for GIN, which not every cluster has, so a missing package explains what it needs instead of failing deep inside buffer construction.DeepEP replaces more than the two collectives. It takes tokens before top-k expansion and replicates them itself, groups arrivals by expert for the grouped GEMM, and reduces the weighted sum in its combine, so the expansion and reduction around the collectives are replaced as well. Its backward pass has no separate entry points: the gradient of a combine is a dispatch and the gradient of a dispatch is a combine, both replayed against the handle the forward dispatch produced.
Limitations
The DeepEP path is disabled with a warning when folded tensor parallelism is active (
tp_size > 1). Folded TP partitions assignments across lanes and restores them by assignment metadata, which a transport whose combine returns token-major rows cannot satisfy. Such layers fall back to the NCCL path rather than silently producing wrong output.Testing
13 unit tests cover backend selection (default, empty, unknown, SM budget), the preflight checks, and gradient shape conformance. Correctness was verified on 16 H100s across two nodes for every routing scenario above, including a forward-and-backward round trip through all four payload collectives.
Follow-up work, kept out of this PR so each change stands on its own: communication/compute overlap (measured 1.11x, composes with this) and deduplication on the NCCL path (1.29x, needs no new dependency).