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 14 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".
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>
Routing weights now leave the dispatch as an autograd output, so the copy cached on the exchange has no readers. Reading weights from there instead of from the graph is what stopped the router gate receiving gradients, so leaving the field in place would invite the same mistake again. Signed-off-by: yh0903 <helloyu0903@gmail.com>
| reduces, which is the same weighted sum the NCCL path performs | ||
| separately after its combine. | ||
| """ | ||
| combined, _, _ = self.buffer.combine(rows, handle=handle, topk_weights=topk_weights, num_sms=self.num_sms) |
There was a problem hiding this comment.
ElasticBuffer.combine(..., topk_weights=...) does not multiply rows by those weights (ElasticBuffer.combine). DeepEP reduces the activations into combined_x and returns the transported weights separately as combined_topk_weights.
Can you apply each aligned routing weight to its expert-output row before the activation reduction, or otherwise implement an equivalent weighted reduction?
There was a problem hiding this comment.
I think this case should be caught by a cosim test -- compare results between DeepEP result and all-to-all result, they should be the same or close. Is there a cosim test in this PR? @yh0903
There was a problem hiding this comment.
Thanks for catching this! You're right, I misunderstood what topk_weights does in ElasticBuffer.combine. It only transports and reduces the weights, but doesn't apply them to rows. I changed _deepep_route to apply the routing weights explicitly before or after the expert MLP based on score_apply, and combine() no longer takes them. I also added a comm-vs-DeepEP parity test for outputs and gradients, plus a check that the router gate gradient is nonzero.
| # which is not known until a batch arrives, and rebuilt if a later | ||
| # batch is larger: variable sequence lengths or a small warm-up batch | ||
| # would otherwise exceed a capacity fixed by the first call. | ||
| if self._deepep_exchange is None or tokens.shape[0] > self._deepep_exchange.num_max_tokens_per_rank: |
There was a problem hiding this comment.
This fixes the original fixed-first-batch capacity issue, but the decision is rank-local. DeepEP requires num_max_tokens_per_rank to have the same value on every rank.
There are two approaches: 1) allreduce (max) the required capacity, and 2) allocate the buffer for the worst case on all ranks.
This is a memory - performance tradeoff. Do you have any idea which is better? We could refer to other frameworks.
There was a problem hiding this comment.
Good point. I ended up doing one all_reduce(MAX) when each layer builds its buffer, so every rank uses the same capacity. With the default comm_max_tokens_per_rank=0, it uses the group-wide first-batch maximum rounded to 512 tokens. For jobs where later batches can be larger, the new comm_max_tokens_per_rank config can provide the worst case explicitly. The buffer isn't resized after initialization; an oversized later batch now gives a clear error instead of making a rank-local resize decision.
| # would otherwise exceed a capacity fixed by the first call. | ||
| if self._deepep_exchange is None or tokens.shape[0] > self._deepep_exchange.num_max_tokens_per_rank: | ||
| if self._deepep_exchange is not None: | ||
| self._deepep_exchange.destroy() |
There was a problem hiding this comment.
In the following pattern, we will use a destroyed exchange.
forward(b1)
forward(b2) # extend buffer
backward(b2)
backward(b1) # Use destroyed exchangeThis actually happens in this use case:
out1 = model(b1)
out2 = model(b2)
loss = out1.sum() + out2.sum()
loss.backward()I'm not sure what we could do for this case. If there is no good solution, we can just detect this pattern and throw an error.
There was a problem hiding this comment.
Yes, this pattern makes runtime resizing unsafe. I removed the resize path entirely. Each layer now builds one exchange and keeps it alive until engine teardown, so both forwards can still use the exchange saved with their handles during backward. If a later batch is too large, it fails before dispatch and points to comm_max_tokens_per_rank instead of destroying and replacing the exchange.
| # an NVLink domain and an RDMA domain rather than assuming a single | ||
| # flat NVLink domain. | ||
| allow_hybrid_mode=True, | ||
| explicitly_destroy=True, |
There was a problem hiding this comment.
With explicitly_destroy=True, the caller needs to explicitly destroy this buffer. The only current call to DeepEPExchange.destroy() is the resize path.
There was a problem hiding this comment.
Yes, I added a registry for the live exchanges, and DeepSpeedEngine.destroy() now destroys all of them in construction order. The engine destructor also calls the same cleanup as a fallback. Since the resize path is gone, buffers are no longer destroyed in the middle of training.
| num_max_tokens_per_rank=tokens.shape[0], | ||
| ) | ||
|
|
||
| received, recv_weights, exchange = deepep_dispatch(self._deepep_exchange, tokens, ro.selected_experts, |
There was a problem hiding this comment.
DeepEP's dispatch kernel doesn't work with FP16, but the code can reach here. Can we reject FP16 settings?
There was a problem hiding this comment.
Yes, added a dtype check at the start of _deepep_route. FP16 now fails before buffer construction or dispatch, with an error pointing users to bfloat16 or comm_backend="comm".
|
|
||
| from deepspeed.utils import logger | ||
|
|
||
| _BACKEND_ENV = "DEEPSPEED_AUTOEP_COMM_BACKEND" |
There was a problem hiding this comment.
Consider move AutoEP comm backend and num SMS into AutoEP configuration. i.e.
comm_backend: Literal["comm", "deepep"] # comm means using comm collective
comm_num_sm: int = 12
There was a problem hiding this comment.
Done. Backend selection is no longer environment-variable based. comm_backend and comm_num_sm now live in the expert_parallel config alongside the rest of AutoEP, with defaults "comm" and 12. validate_autoep_config() rejects unknown backends and non-positive SM budgets.
| smaller than DeepEP's automatic choice, which assumes it is the only thing | ||
| on the fabric. | ||
| """ | ||
| return num_sms + 4 |
There was a problem hiding this comment.
could num qps also be configurable? Could be called qp_margin and make it default to 4
There was a problem hiding this comment.
Yes, the new expert_parallel.comm_qp_margin field defaults to 4 and is validated as a non-negative integer. The buffer reserves comm_num_sm + comm_qp_margin queue pairs.
| Used as the backward of a combine, where the weights matter as much as | ||
| the rows: their gradient is what reaches the router gate. | ||
| """ | ||
| recv_x, _, recv_weights, _, _ = self.buffer.dispatch( |
There was a problem hiding this comment.
both dispatch and combine support passing num_qps, however here only dispatch has this optional parameter. If num_qps is the same as where buffer is initialized, why not omit it in dispatch call? We can avoid store the num_qps in the class and call self.buffer.get_theoretical_num_qps(self.num_sms)
There was a problem hiding this comment.
Fixed. The QP count is derived once from num_sms + qp_margin and passed as num_allocated_qps when the ElasticBuffer is constructed. It is no longer stored as a separate DeepEPExchange field, and neither dispatch() nor combine() passes a per-call num_qps; those calls only pass num_sms.
Buffer construction is collective, but the capacity came from the local token count, which differs per rank. Some ranks entered construction while others did not, and those that did waited for peers that never arrived until the connection dropped. Reducing the count first makes the capacity identical everywhere, which also makes the rebuild decision unanimous without a second collective. Routing weights arrive in a worst-case buffer like the rows do, but only the rows were trimmed to the arrivals the prefix sum reports. Combine then received a row count and a weight count that disagreed. This was introduced when the weights started coming back through autograd rather than being read off the exchange already trimmed. Signed-off-by: yh0903 <helloyu0903@gmail.com>
| # them by assignment metadata, which a transport whose combine | ||
| # returns token-major rows cannot satisfy. Saying so beats | ||
| # silently running something other than what was asked for. | ||
| logger.warning("AutoEP: the DeepEP backend does not support folded tensor parallelism; " |
There was a problem hiding this comment.
I saw a behavior incosistency here. When DeepEP cannot be imported, it will abort. When DeepEP does not compatible with TP folding, it will fallback. I'll favor both use abort because fallback will make user/agent think DeepEP is used even if it fallback to NCCL, if user are not careful enough to look at the log.
| # keeps the two consistent. | ||
| prefix = handle.psum_num_recv_tokens_per_expert | ||
| counts = torch.diff(prefix, prepend=prefix.new_zeros(1)).to(torch.int32) | ||
| received = received[:int(prefix[-1].item())] |
There was a problem hiding this comment.
Non-blocking nit: unnecessary D2H sync from .item() in the dispatch hot path
received[:int(prefix[-1].item())] forces a GPU→CPU synchronization before the expert GEMM. The .item() drains the GPU pipeline (~20-50μs per layer) because the CPU must wait for the dispatch kernel to finish, transfer the scalar over PCIe, and only then can enqueue the GEMM.
This trim appears unnecessary: all three GroupedExperts paths handle the valid row range internally — torch._grouped_mm and the Triton kernel use offsets = cumsum(counts) to define group boundaries (rows beyond offsets[-1] are not part of any group), and the for-loop path does its own x[:sum(counts)] trim. So padding rows in the untrimmed received buffer should be safely ignored.
Consider removing the trim and passing the full received buffer directly. Estimated impact: ~0.5% of step time. Minor, but the fix is zero-cost.
|
Hi @yh0903 , very glad to see DeepEP being integrated into DeepSpeed. I have left my comments. Thanks! |
Correctness: DeepEP's combine does not multiply rows by the topk_weights it is handed. It transports and reduces them separately and returns them as a second output, so passing the routing weights there dropped them from the result and the expert outputs came back summed but unweighted. The layer now applies them itself, before or after the experts to match what score_apply means on the collective path, and never passes them to combine. Configuration: comm_backend, comm_num_sm and comm_qp_margin move from environment variables into the expert_parallel config section, alongside the rest of AutoEP. The default names the transport rather than the library, since deepspeed.comm is not NCCL on every accelerator. Buffer lifetime: a buffer outgrown by a larger batch is retained rather than destroyed, because a backward from an earlier forward replays against the buffer that forward dispatched on. Capacity is rounded up so a slowly growing sequence length does not rebuild on nearly every step, and engine.destroy() releases every buffer, which nothing did before. Also reject fp16, which DeepEP's dispatch kernel cannot handle; refuse folded tensor parallelism rather than silently falling back to the collective path; drop the device-to-host sync that trimmed rows the expert paths already bound by their counts; and stop passing a queue pair count that the buffer derives from the SM count itself. Signed-off-by: yh0903 <helloyu0903@gmail.com>
Removing the trim in front of the expert GEMM also removed it in front of combine, which is not safe: combine reads the rows the handle describes, so handing it a worst-case buffer reads past them. That does not raise, it faults, and every rank died without writing a traceback. The count now comes from handle.num_expanded_tokens, which is already a Python int, so the rows are cut without the device-to-host synchronisation that reading the end of the prefix sum would have put in front of every layer's GEMM. Signed-off-by: yh0903 <helloyu0903@gmail.com>
The agreement is collective and ends in a device-to-host read, so running it on every layer of every step put a synchronisation in the middle of the forward pass. Measured end to end, that alone turned a 1.16x speedup into a 0.84x slowdown. Ranks in an expert-parallel group are handed the same micro-batch shape, so caching the agreement against the shape that produced it keeps the decision to skip it unanimous, and the collective still runs whenever the shape changes, which is when a rebuild might be needed. Documentation now carries the step times measured with the routing weights actually applied: 360.6 ms to 310.6 ms, a 1.16x speedup. Signed-off-by: yh0903 <helloyu0903@gmail.com>
acf2e5b to
ab4e7cc
Compare
| # micro-batch shape, so they leave and re-enter the agreement together, | ||
| # and a collective on every layer of every step would otherwise put a | ||
| # device-to-host synchronisation in the middle of the forward pass. | ||
| if tokens.shape[0] != self._deepep_tokens_per_rank: |
There was a problem hiding this comment.
I think there still have chance that some ranks meet the condition and some ranks didn't meet this condition. Thus some ranks may execute allreduce and some are not, leading to a stall.
How much is the allreduce cost? We may need to live with this cost if we want to use memory smartly.
The recorded sweep predated the fix that made the routing weights actually reach the reduction, so it described a step that was doing less work than the one users get. Re-measuring also required fixing how the comparison was run: the two backends had been timed in separate jobs, on separate pods, and the collective baseline alone varied by a quarter between jobs, which is far larger than the differences being compared. Both backends are now timed in one job on one set of pods, rotated through every position so no arm keeps the same slot, and reported as the fastest of three runs -- interference only ever makes a step slower, so the fastest observation is the least contaminated one. A mean of two runs had put 8 SMs ahead of 12 until one disturbed run was identified as the whole of the difference. The default of 12 is unchanged and now has measurements behind it: 325.1 ms on the collective path against 256.9 ms on DeepEP, a 1.27x speedup. Signed-off-by: yh0903 <helloyu0903@gmail.com>
The fastest of several runs is the least contaminated estimate of what the code can do, but a training job gets the typical run, so a headline speedup taken from the best observation overstates what a user will see. The reported figures are now medians of three: 329.4 ms on the collective path against 265.4 ms on DeepEP, a 1.24x speedup rather than the 1.27x the fastest runs gave. Spread is quoted with them because it is around 20% run to run, which is wider than the gap between neighbouring SM budgets and is the reason several earlier single-run comparisons disagreed with each other. Signed-off-by: yh0903 <helloyu0903@gmail.com>
…ive path Resizing on demand made the agreement collective but conditional on the local token count. Ranks that had grown would enter the all-reduce and ranks that had not would skip it, which hangs rather than fails, and justifying the skip on ranks seeing identical micro-batch shapes only holds until they do not. Agreeing every time instead costs a device-to-host read in the middle of every layer of every step, measured at roughly a fifth of the step. The size is now agreed once, on the first forward, which every rank reaches together, and comm_max_tokens_per_rank lets a job whose later batches are larger say so up front. A batch that outgrows the buffer raises and names that setting. The parity test compares a step through each transport, activations and gradients, and asserts the router gate's gradient is present and nonzero. This is the shape of test that would have caught the routing weights being handed to a combine that transports but does not apply them: the loss still fell and every mock-level test still passed. Signed-off-by: yh0903 <helloyu0903@gmail.com>
Two independent jobs gave 1.24x and 1.21x. Quoting the higher one takes the better of two measurements when both are equally valid, so the figure is now the pooled one with both quoted alongside it. Worth recording which side the variation comes from: DeepEP's own median moved 0.3% between the jobs while the collective baseline moved 2.5%, so the spread in the ratio is the baseline's, not the transport's. Signed-off-by: yh0903 <helloyu0903@gmail.com>
|
@yh0903 thank you for addressing my comments. |
Thanks @delock for the follow-up review and approval. Sorry I didn't response to your message timely, and actually I was typing the replies! Really appreciate your comments again, which led to multiple improvements! I should have replied to your thread to make it easier to track the latest changes, but thanks for looking into them! |
Several comments read like design docs: the SM-budget default, the buffer-sizing rationale, and the routing-weight placement note each re-derived their reasoning in full rather than stating it once. Cut them to the conclusion a maintainer needs at the call site; the fuller argument for the SM default was already in autoep.rst, and the deadlock reasoning for buffer sizing keeps its one sentence rather than the paragraph. No code changes -- verified by comparing the AST before and after. Signed-off-by: yh0903 <helloyu0903@gmail.com>
4 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).