From cbaa06ef5889496e19a3a9c6701a4d89586a49c9 Mon Sep 17 00:00:00 2001 From: yh0903 Date: Tue, 4 Aug 2026 10:39:20 -0700 Subject: [PATCH 1/7] Add an opt-in DeepEP transport for the AutoEP expert all-to-all 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= (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 --- deepspeed/module_inject/auto_ep_comm.py | 293 ++++++++++++++++++ deepspeed/module_inject/auto_ep_layer.py | 83 ++++- tests/unit/module_inject/test_auto_ep_comm.py | 118 +++++++ 3 files changed, 488 insertions(+), 6 deletions(-) create mode 100644 deepspeed/module_inject/auto_ep_comm.py create mode 100644 tests/unit/module_inject/test_auto_ep_comm.py diff --git a/deepspeed/module_inject/auto_ep_comm.py b/deepspeed/module_inject/auto_ep_comm.py new file mode 100644 index 000000000000..5217a50c1289 --- /dev/null +++ b/deepspeed/module_inject/auto_ep_comm.py @@ -0,0 +1,293 @@ +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +"""Selectable transports for the AutoEP expert all-to-all. + +The dispatch and combine collectives are the largest single cost in an AutoEP +step, and a measured replay of real routing on 16 H100s across two nodes put +NCCL at 195.8 ms of payload all-to-all per step against DeepEP's 86.3 ms. This +module is what lets that be switched without the MoE layer knowing which +transport it is using. + +Selection is by environment variable and defaults to the NCCL path, so a job +that sets nothing behaves exactly as before: + + DEEPSPEED_AUTOEP_COMM_BACKEND=nccl (default) + DEEPSPEED_AUTOEP_COMM_BACKEND=deepep + +A backend is asked for once per layer and reused, because DeepEP's buffers are +sized at construction and are expensive to rebuild. +""" + +from __future__ import annotations + +import os + +import torch + +from deepspeed.utils import logger + +_BACKEND_ENV = "DEEPSPEED_AUTOEP_COMM_BACKEND" +_NUM_SMS_ENV = "DEEPSPEED_AUTOEP_COMM_SMS" +NCCL_BACKEND = "nccl" +# Names the library, not its version: v2 is the only path implemented, and +# nothing about the name would have to change if that ever grew. +DEEPEP_BACKEND = "deepep" +AVAILABLE_BACKENDS = (NCCL_BACKEND, DEEPEP_BACKEND) +# GIN, and so DeepEP, does not exist in any form below this NCCL version. +NCCL_GIN_MIN_VERSION = (2, 30, 4) +# Chosen by sweeping whole training steps rather than the collective alone. +# On 16 H100s across two nodes the step took 340, 311, 353, 360 and 391 ms at 8, +# 12, 16, 24 and 32 SMs. Twelve is the point where the collective is already as +# fast as it gets while the expert GEMM still has the SMs it needs: at 8 the +# collective itself degrades, and above 12 the step grows because communication +# takes SMs the rest of the step was using. +DEFAULT_COMM_SMS = 12 + + +def _qps_for_sms(num_sms: int) -> int: + """Queue pairs to reserve for a given SM count. + + One per SM plus a small margin for the control path. This is deliberately + smaller than DeepEP's automatic choice, which assumes it is the only thing + on the fabric. + """ + return num_sms + 4 + + +def configured_backend() -> str: + """The transport this process should use for expert all-to-all. + + An unset variable selects NCCL, so a job that opts into nothing keeps the + behaviour it had before this existed. + """ + raw = os.environ.get(_BACKEND_ENV) + if raw is None or not raw.strip(): + return NCCL_BACKEND + name = raw.strip().lower() + if name not in AVAILABLE_BACKENDS: + raise ValueError(f"{_BACKEND_ENV}={raw!r} is not one of {AVAILABLE_BACKENDS}") + return name + + +def configured_num_sms() -> int: + """SM budget for communication; 0 lets the backend decide. + + Worth setting because the collective competes with the expert GEMM for + SMs: measurements showed the GEMM running 1.21-1.25x slower whenever a + collective was in flight. + """ + raw = os.environ.get(_NUM_SMS_ENV) + if raw is None or not raw.strip(): + return 0 + try: + return int(raw) + except ValueError as error: + raise ValueError(f"{_NUM_SMS_ENV}={raw!r} is not an integer") from error + + +def _import_deep_ep(): + """Import DeepEP, explaining the environment it needs when it is absent. + + DeepEP is an optional dependency with prerequisites a cluster either meets + or does not, and the failures it produces otherwise are opaque: a missing + GIN-capable NCCL surfaces as an assertion inside buffer construction rather + than as anything naming NCCL. Since this backend is only ever reached by + explicit opt-in, the person who opted in is the one who can act on this. + """ + try: + import deep_ep + except ImportError as error: + raise ImportError( + f"{_BACKEND_ENV}={DEEPEP_BACKEND} requires the deep_ep package, which is not installed. It also " + "requires NCCL 2.30.4 or newer built with GIN support: the transport is unavailable below that " + "version regardless of the network. Unset " + f"{_BACKEND_ENV} to use the default NCCL all-to-all, which has no such requirement.") from error + + nccl_version = _nccl_version() + if nccl_version is not None and nccl_version < NCCL_GIN_MIN_VERSION: + installed = ".".join(str(part) for part in nccl_version) + minimum = ".".join(str(part) for part in NCCL_GIN_MIN_VERSION) + # Deliberately a warning. This reports the NCCL that torch bundles and + # loads through its own RPATH, which DeepEP need not be using: DeepEP + # links the NCCL it was built against, and that pairing has been + # observed working while torch reported an older one. Refusing to start + # on this signal would block a configuration already known to run. + logger.warning( + f"torch reports NCCL {installed}, older than the {minimum} that GIN requires. DeepEP links its own " + "NCCL, so this is only a problem if it also resolves to the older one; a failure inside buffer " + f"construction is the symptom. Unset {_BACKEND_ENV} to fall back to the default all-to-all.") + return deep_ep + + +def _nccl_version() -> tuple[int, ...] | None: + """The NCCL version torch is linked against, or None if unknowable.""" + try: + return tuple(torch.cuda.nccl.version()) #ignore-cuda + except Exception: + # Not being able to tell is not a reason to block a run that might work. + return None + + +class DeepEPExchange: + """Wraps a DeepEP v2 ``ElasticBuffer`` for one MoE layer. + + Only v2 is supported. The legacy v1 ``Buffer`` moves data over NVSHMEM and + IBGDA instead of NCCL, which needs either the NVreg_EnableStreamMemOPs + driver parameter or the GDRCopy device, and it reports markedly lower + internode bandwidth -- the case this backend exists to improve. + + DeepEP has no separate backward 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, so the handle has to + survive from forward to backward. + """ + + def __init__(self, ep_group, num_experts: int, top_k: int, hidden_size: int, num_max_tokens_per_rank: int): + deep_ep = _import_deep_ep() + + self.deep_ep = deep_ep + # Queue pairs are the scarce resource here. Left automatic, DeepEP + # claims 65 to 129 of them, which is fine in a process that does + # nothing else but fails in a training step where ZeRO and the + # data-parallel groups have already taken their share. Asking for only + # what the chosen SM count needs keeps the request proportionate. + num_sms = configured_num_sms() or DEFAULT_COMM_SMS + self.buffer = deep_ep.ElasticBuffer( + ep_group, + num_max_tokens_per_rank=num_max_tokens_per_rank, + hidden=hidden_size, + num_topk=top_k, + use_fp8_dispatch=False, + num_allocated_qps=_qps_for_sms(num_sms), + # Required once the EP group spans nodes: it splits the ranks into + # an NVLink domain and an RDMA domain rather than assuming a single + # flat NVLink domain. + allow_hybrid_mode=True, + explicitly_destroy=True, + ) + self.num_sms = num_sms + self.num_qps = self.buffer.get_theoretical_num_qps(self.num_sms) + self.num_experts = num_experts + # The handle the last dispatch produced. Combine and both backward + # passes replay against it, so it has to outlive the dispatch call. + self.last_handle = None + # The routing weights that arrived with the last dispatch, kept so the + # layer can decide whether to apply them before or after the experts. + self.last_recv_weights = None + + def dispatch(self, tokens: torch.Tensor, topk_idx: torch.Tensor, topk_weights: torch.Tensor): + """Send tokens to their experts, returning rows, weights and handle. + + The weights travel with the tokens because the reduction that uses + them happens on the receiving side, after the experts have run. + """ + recv_x, _, recv_weights, handle, _ = self.buffer.dispatch( + tokens, + topk_idx=topk_idx.to(self.deep_ep.topk_idx_t), + # DeepEP reduces in float32, and the router's scores may be bf16. + topk_weights=topk_weights.float(), + num_experts=self.num_experts, + # Group arrivals by expert rather than by source rank. The + # grouped GEMM walks contiguous per-expert ranges, so the default + # source-major layout has the right number of rows in an order the + # GEMM cannot use. + do_expand=True, + # No per-expert padding: the counts that become the GEMM's group + # offsets have to describe the rows that are actually there. + expert_alignment=1, + num_sms=self.num_sms, + num_qps=self.num_qps, + ) + # The returned event only holds anything when the call was made with + # async_with_compute_stream; a synchronous result is already usable. + self.last_handle = handle + self.last_recv_weights = recv_weights + return recv_x, recv_weights, handle + + def dispatch_with_handle(self, tokens: torch.Tensor, handle) -> torch.Tensor: + recv_x, _, _, _, _ = self.buffer.dispatch( + tokens, + handle=handle, + num_sms=self.num_sms, + num_qps=self.num_qps, + ) + return recv_x + + def combine(self, rows: torch.Tensor, handle, topk_weights=None) -> torch.Tensor: + """Reduce expert outputs back to their tokens. + + Passing ``topk_weights`` makes DeepEP weight each expert's output as it + 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) + return combined + + def destroy(self) -> None: + self.buffer.destroy() + + +def _conform_rows(tensor: torch.Tensor, shape) -> torch.Tensor: + """Trim or zero-extend ``tensor`` to ``shape``'s row count. + + DeepEP returns whole buffers sized for the worst case, but autograd checks + a gradient against the exact input it corresponds to. Rows beyond the ones + that carried tokens hold no gradient, so trimming discards nothing and + extending contributes nothing. + """ + rows = shape[0] + if tensor.shape[0] == rows: + return tensor + if tensor.shape[0] > rows: + return tensor[:rows] + extended = tensor.new_zeros((rows, tensor.shape[1])) + extended[:tensor.shape[0]] = tensor + return extended + + +class _DeepEPDispatch(torch.autograd.Function): + """Forward dispatch whose backward is the matching combine.""" + + @staticmethod + def forward(ctx, exchange: DeepEPExchange, tokens: torch.Tensor, topk_idx: torch.Tensor, + topk_weights: torch.Tensor): + received, recv_weights, handle = exchange.dispatch(tokens, topk_idx, topk_weights) + ctx.exchange = exchange + ctx.handle = handle + ctx.tokens_shape = tokens.shape + ctx.recv_weights = recv_weights + return received + + @staticmethod + 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 + + +class _DeepEPCombine(torch.autograd.Function): + """Combine whose backward is the matching dispatch, on the same handle.""" + + @staticmethod + def forward(ctx, exchange: DeepEPExchange, rows: torch.Tensor, handle, topk_weights): + ctx.exchange = exchange + ctx.handle = handle + # The forward input was trimmed to the rows that actually arrived, + # while the backward dispatch hands back a whole buffer. Autograd + # requires the gradient to match the input it is the gradient of. + ctx.rows_shape = rows.shape + return exchange.combine(rows, handle, topk_weights) + + @staticmethod + def backward(ctx, grad_combined): + grad_rows = ctx.exchange.dispatch_with_handle(grad_combined.contiguous(), ctx.handle) + return None, _conform_rows(grad_rows, ctx.rows_shape), None, None + + +def deepep_dispatch(exchange: DeepEPExchange, tokens: torch.Tensor, topk_idx: torch.Tensor, + topk_weights: torch.Tensor): + received = _DeepEPDispatch.apply(exchange, tokens, topk_idx, topk_weights) + return received, exchange + + +def deepep_combine(exchange: DeepEPExchange, rows: torch.Tensor, handle, topk_weights=None) -> torch.Tensor: + return _DeepEPCombine.apply(exchange, rows, handle, topk_weights) diff --git a/deepspeed/module_inject/auto_ep_layer.py b/deepspeed/module_inject/auto_ep_layer.py index 7a2579017b4a..a377de8adb20 100644 --- a/deepspeed/module_inject/auto_ep_layer.py +++ b/deepspeed/module_inject/auto_ep_layer.py @@ -22,6 +22,8 @@ from deepspeed.module_inject.auto_ep_config import AutoEPConfig, MoELayerSpec, resolve_autoep_config_defaults from deepspeed.module_inject.auto_ep_folding import mark_autoep_folding_router_parameter from deepspeed.utils import logger +from deepspeed.module_inject.auto_ep_comm import (DEEPEP_BACKEND, NCCL_BACKEND, DeepEPExchange, configured_backend, + deepep_combine, deepep_dispatch) from deepspeed.moe.ep_router import TokenChoiceTopKRouter from deepspeed.moe.ep_count import count_tokens_per_expert from deepspeed.moe.ep_experts import GroupedExperts @@ -519,6 +521,15 @@ def __init__( # Router-logit cache self._cached_router_logits = None + # Resolved once per layer: a DeepEP exchange sizes its buffer at + # construction, so it is built on first use and kept. A bad value must + # not break the default path, which nobody opted into. + try: + self.comm_backend = configured_backend() + except ValueError as error: + logger.warning(f"AutoEP: falling back to the default all-to-all: {error}") + self.comm_backend = NCCL_BACKEND + self._deepep_exchange = None self._register_logit_hook() def _register_logit_hook(self): @@ -571,6 +582,60 @@ def set_deepspeed_parallelism( ) self.ep_group = groups._get_expert_parallel_group(self.ep_group_name) + def _deepep_route(self, tokens: torch.Tensor, ro: "RouterOutput") -> torch.Tensor: + """Dispatch, run the experts, and combine through DeepEP. + + ``tokens`` is [T, H] before top-k expansion. DeepEP replicates each + token to the ranks that need it rather than being handed one row per + selected expert, groups arrivals by expert for the grouped GEMM, and + reduces the weighted sum on the way back. It therefore replaces the + expansion and the reduction around the collectives, not just the + collectives, and returns [T, H] ready for the shared tail of forward. + """ + if self._deepep_exchange is None: + # Built on first use rather than at construction, because the + # buffer is sized from the token count and that is not known until + # a batch arrives. + self._deepep_exchange = DeepEPExchange( + ep_group=self.ep_group, + num_experts=self.num_experts, + top_k=self.top_k, + hidden_size=self.hidden_size, + num_max_tokens_per_rank=tokens.shape[0], + ) + + received, exchange = deepep_dispatch(self._deepep_exchange, tokens, ro.selected_experts, ro.top_scores) + handle = exchange.last_handle + + # Two quantities have to line up for the grouped GEMM, and they come + # from different places: the rows are the valid prefix of a buffer + # sized for the worst case, and the per-expert counts arrive as a + # prefix sum. Differencing the same prefix sum that bounds the rows + # 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())] + if counts.numel() != self.num_local_experts: + raise RuntimeError(f"DeepEP returned {counts.numel()} expert counts, but this rank owns " + f"{self.num_local_experts} experts") + + # The router weights must be applied exactly once. In "pre" mode the + # NCCL path scales tokens before the experts see them, so the arrived + # rows are scaled the same way; in "post" mode the scaling is part of + # the reduction and is handed to combine instead. In expand mode each + # row is one token-expert slot, so a weight per row broadcasts across + # the hidden dimension. + recv_weights = exchange.last_recv_weights + if self.score_apply == "pre" and recv_weights is not None: + scale = recv_weights[:received.shape[0]].reshape(-1, 1) + received = (received.to(torch.float32) * scale).to(received.dtype) + combine_weights = None + else: + combine_weights = recv_weights + + expert_output = self.experts(received, counts) + return deepep_combine(exchange, expert_output, handle, combine_weights) + def forward( self, hidden_states: torch.Tensor, @@ -679,14 +744,17 @@ def forward( ep_group=self.ep_group, ) - routed_input = _AllToAllV.apply(self.ep_group, routed_input, plan.input_splits, plan.output_splits) + if self.comm_backend == DEEPEP_BACKEND: + expert_output = self._deepep_route(x, ro) + else: + routed_input = _AllToAllV.apply(self.ep_group, routed_input, plan.input_splits, plan.output_splits) - routed_input, perm_indices, aligned_counts, n_tokens = permute_by_local_expert( - routed_input, plan.local_counts_by_source) - expert_output = self.experts(routed_input, aligned_counts) - expert_output = unpermute_by_local_expert(expert_output, perm_indices, n_tokens) + routed_input, perm_indices, aligned_counts, n_tokens = permute_by_local_expert( + routed_input, plan.local_counts_by_source) + expert_output = self.experts(routed_input, aligned_counts) + expert_output = unpermute_by_local_expert(expert_output, perm_indices, n_tokens) - expert_output = _AllToAllV.apply(self.ep_group, expert_output, plan.output_splits, plan.input_splits) + expert_output = _AllToAllV.apply(self.ep_group, expert_output, plan.output_splits, plan.input_splits) if folded_tp: output = restore_combined(expert_output, @@ -694,6 +762,9 @@ def forward( tp_group=self.tp_group, validate_coverage=self.validate_folding_routing).reshape(bsz, seqlen, hdim) self._last_folding_dispatch_counters = dispatch_counters(restore_ctx) + elif self.comm_backend == DEEPEP_BACKEND: + # Already reduced over top-k and back in token order. + output = expert_output.reshape(bsz, seqlen, hdim) else: output = combine_from_routed( expert_output, diff --git a/tests/unit/module_inject/test_auto_ep_comm.py b/tests/unit/module_inject/test_auto_ep_comm.py new file mode 100644 index 000000000000..6d4904e675a7 --- /dev/null +++ b/tests/unit/module_inject/test_auto_ep_comm.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team + +import os +import sys +import unittest + +import torch +from unittest import mock + +from deepspeed.module_inject.auto_ep_comm import (AVAILABLE_BACKENDS, DEEPEP_BACKEND, NCCL_BACKEND, _conform_rows, + _import_deep_ep, configured_backend, configured_num_sms) + + +class TestAutoEPCommBackendSelection(unittest.TestCase): + + def test_defaults_to_nccl(self): + # An unset variable must leave existing jobs on the shipped path. + with mock.patch.dict(os.environ, {}, clear=True): + self.assertEqual(configured_backend(), NCCL_BACKEND) + self.assertEqual(configured_num_sms(), 0) + + def test_selects_deepep(self): + with mock.patch.dict(os.environ, {"DEEPSPEED_AUTOEP_COMM_BACKEND": "DeepEP"}): + self.assertEqual(configured_backend(), DEEPEP_BACKEND) + + def test_rejects_unknown_backend(self): + # Failing loudly beats silently running the wrong transport. + with mock.patch.dict(os.environ, {"DEEPSPEED_AUTOEP_COMM_BACKEND": "moonep"}): + with self.assertRaises(ValueError) as caught: + configured_backend() + self.assertIn(str(AVAILABLE_BACKENDS), str(caught.exception)) + + def test_sm_budget_is_configurable(self): + with mock.patch.dict(os.environ, {"DEEPSPEED_AUTOEP_COMM_SMS": "24"}): + self.assertEqual(configured_num_sms(), 24) + + def test_blank_value_is_treated_as_unset(self): + # An empty variable is what an unset shell variable expands to. + with mock.patch.dict(os.environ, {"DEEPSPEED_AUTOEP_COMM_BACKEND": " "}): + self.assertEqual(configured_backend(), NCCL_BACKEND) + with mock.patch.dict(os.environ, {"DEEPSPEED_AUTOEP_COMM_SMS": ""}): + self.assertEqual(configured_num_sms(), 0) + + def test_non_numeric_sm_budget_is_rejected(self): + with mock.patch.dict(os.environ, {"DEEPSPEED_AUTOEP_COMM_SMS": "many"}): + with self.assertRaises(ValueError): + configured_num_sms() + + +class TestDeepEPPreflight(unittest.TestCase): + """Opting in on an unsuitable machine should say what is missing.""" + + def test_missing_package_names_its_requirements(self): + with mock.patch.dict(sys.modules, {"deep_ep": None}): + with self.assertRaises(ImportError) as caught: + _import_deep_ep() + message = str(caught.exception) + self.assertIn("2.30.4", message) + self.assertIn("DEEPSPEED_AUTOEP_COMM_BACKEND", message) + + def test_old_torch_nccl_warns_but_does_not_block(self): + # torch reports the NCCL it bundles, which DeepEP need not be using; + # DeepEP has been measured working while torch reported 2.28.9, so this + # signal must not stop a run. + module = mock.MagicMock() + with mock.patch.dict(sys.modules, {"deep_ep": module}): + with mock.patch("deepspeed.module_inject.auto_ep_comm._nccl_version", return_value=(2, 28, 9)): + with mock.patch("deepspeed.module_inject.auto_ep_comm.logger") as log: + self.assertIs(_import_deep_ep(), module) + message = log.warning.call_args[0][0] + self.assertIn("2.28.9", message) + self.assertIn("2.30.4", message) + + def test_new_enough_nccl_passes(self): + module = mock.MagicMock() + with mock.patch.dict(sys.modules, {"deep_ep": module}): + with mock.patch("deepspeed.module_inject.auto_ep_comm._nccl_version", return_value=(2, 30, 4)): + self.assertIs(_import_deep_ep(), module) + + def test_unknown_nccl_version_does_not_block(self): + # Failing to detect a version is not evidence of an unusable one. + module = mock.MagicMock() + with mock.patch.dict(sys.modules, {"deep_ep": module}): + with mock.patch("deepspeed.module_inject.auto_ep_comm._nccl_version", return_value=None): + self.assertIs(_import_deep_ep(), module) + + +class TestGradientConformance(unittest.TestCase): + """Autograd checks a gradient against the exact input it belongs to.""" + + def test_trims_a_longer_buffer(self): + # DeepEP returns buffers sized for the worst case; the rows past the + # ones that carried tokens hold no gradient. + grad = torch.ones((10, 4)) + + conformed = _conform_rows(grad, (6, 4)) + + self.assertEqual(tuple(conformed.shape), (6, 4)) + self.assertTrue(torch.equal(conformed, torch.ones((6, 4)))) + + def test_extends_a_shorter_buffer_with_zeros(self): + grad = torch.ones((3, 4)) + + conformed = _conform_rows(grad, (5, 4)) + + self.assertEqual(tuple(conformed.shape), (5, 4)) + self.assertTrue(torch.equal(conformed[:3], torch.ones((3, 4)))) + self.assertTrue(torch.equal(conformed[3:], torch.zeros((2, 4)))) + + def test_matching_shape_is_passed_through_untouched(self): + grad = torch.randn((7, 4)) + + self.assertIs(_conform_rows(grad, (7, 4)), grad) + + +if __name__ == "__main__": + unittest.main() From 79a6ea0b19df2a97f89337af991d4a5aa9690d90 Mon Sep 17 00:00:00 2001 From: yh0903 Date: Tue, 4 Aug 2026 11:17:40 -0700 Subject: [PATCH 2/7] Address review feedback on the DeepEP backend 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 --- deepspeed/module_inject/auto_ep_comm.py | 59 +++++++++++++++---- deepspeed/module_inject/auto_ep_layer.py | 34 ++++++++--- docs/code-docs/source/autoep.rst | 31 ++++++++++ tests/unit/module_inject/test_auto_ep_comm.py | 48 ++++++++++++++- 4 files changed, 151 insertions(+), 21 deletions(-) diff --git a/deepspeed/module_inject/auto_ep_comm.py b/deepspeed/module_inject/auto_ep_comm.py index 5217a50c1289..5b2d03a05c8c 100644 --- a/deepspeed/module_inject/auto_ep_comm.py +++ b/deepspeed/module_inject/auto_ep_comm.py @@ -168,6 +168,8 @@ def __init__(self, ep_group, num_experts: int, top_k: int, hidden_size: int, num self.num_sms = num_sms self.num_qps = self.buffer.get_theoretical_num_qps(self.num_sms) self.num_experts = num_experts + # Recorded so the layer can tell when a later batch outgrows it. + self.num_max_tokens_per_rank = num_max_tokens_per_rank # The handle the last dispatch produced. Combine and both backward # passes replay against it, so it has to outlive the dispatch call. self.last_handle = None @@ -204,14 +206,31 @@ def dispatch(self, tokens: torch.Tensor, topk_idx: torch.Tensor, topk_weights: t self.last_recv_weights = recv_weights return recv_x, recv_weights, handle - def dispatch_with_handle(self, tokens: torch.Tensor, handle) -> torch.Tensor: - recv_x, _, _, _, _ = self.buffer.dispatch( + def dispatch_with_handle(self, tokens: torch.Tensor, handle): + """Replay a dispatch against a cached handle, keeping both outputs. + + 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( tokens, handle=handle, num_sms=self.num_sms, num_qps=self.num_qps, ) - return recv_x + return recv_x, recv_weights + + def combine_with_weight_grad(self, rows: torch.Tensor, handle, weight_grads=None): + """Combine that also returns the reduced weight gradient. + + Used as the backward of a dispatch. The weight gradient travels the + same path as the row gradient, so it is reduced by the same call. + """ + combined, combined_weights, _ = self.buffer.combine(rows, + handle=handle, + topk_weights=weight_grads, + num_sms=self.num_sms) + return combined, combined_weights def combine(self, rows: torch.Tensor, handle, topk_weights=None) -> torch.Tensor: """Reduce expert outputs back to their tokens. @@ -255,13 +274,24 @@ def forward(ctx, exchange: DeepEPExchange, tokens: torch.Tensor, topk_idx: torch ctx.exchange = exchange ctx.handle = handle ctx.tokens_shape = tokens.shape - ctx.recv_weights = recv_weights - return received + ctx.weights_shape = None if topk_weights is None else topk_weights.shape + # Dispatch moves the weights alongside the tokens, so the received + # copies are what downstream code differentiates; returning them makes + # autograd carry their gradient back to the router gate. Without this + # the gate silently receives nothing and stops learning. + return received, recv_weights @staticmethod - 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 + def backward(ctx, grad_received, grad_recv_weights): + grad_tokens, grad_weights = ctx.exchange.combine_with_weight_grad( + grad_received.contiguous(), + ctx.handle, + None if grad_recv_weights is None else grad_recv_weights.contiguous(), + ) + conformed_weights = None + if grad_weights is not None and ctx.weights_shape is not None: + conformed_weights = grad_weights[:ctx.weights_shape[0]].reshape(ctx.weights_shape) + return None, _conform_rows(grad_tokens, ctx.tokens_shape), None, conformed_weights class _DeepEPCombine(torch.autograd.Function): @@ -275,18 +305,23 @@ def forward(ctx, exchange: DeepEPExchange, rows: torch.Tensor, handle, topk_weig # while the backward dispatch hands back a whole buffer. Autograd # requires the gradient to match the input it is the gradient of. ctx.rows_shape = rows.shape + ctx.weights_shape = None if topk_weights is None else topk_weights.shape return exchange.combine(rows, handle, topk_weights) @staticmethod def backward(ctx, grad_combined): - grad_rows = ctx.exchange.dispatch_with_handle(grad_combined.contiguous(), ctx.handle) - return None, _conform_rows(grad_rows, ctx.rows_shape), None, None + grad_rows, grad_weights = ctx.exchange.dispatch_with_handle(grad_combined.contiguous(), ctx.handle) + conformed_weights = None + if grad_weights is not None and ctx.weights_shape is not None: + conformed_weights = grad_weights[:ctx.weights_shape[0]] + return None, _conform_rows(grad_rows, ctx.rows_shape), None, conformed_weights def deepep_dispatch(exchange: DeepEPExchange, tokens: torch.Tensor, topk_idx: torch.Tensor, topk_weights: torch.Tensor): - received = _DeepEPDispatch.apply(exchange, tokens, topk_idx, topk_weights) - return received, exchange + """Dispatch tokens and their routing weights, keeping both differentiable.""" + received, recv_weights = _DeepEPDispatch.apply(exchange, tokens, topk_idx, topk_weights) + return received, recv_weights, exchange def deepep_combine(exchange: DeepEPExchange, rows: torch.Tensor, handle, topk_weights=None) -> torch.Tensor: diff --git a/deepspeed/module_inject/auto_ep_layer.py b/deepspeed/module_inject/auto_ep_layer.py index 36fa0f06e8ba..3dbfcc636129 100644 --- a/deepspeed/module_inject/auto_ep_layer.py +++ b/deepspeed/module_inject/auto_ep_layer.py @@ -561,6 +561,14 @@ def set_deepspeed_parallelism( if folding_group_handles is not None: self.folding_group_handles = folding_group_handles + if self.comm_backend == DEEPEP_BACKEND and folding_group_handles.spec.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. Saying so beats + # silently running something other than what was asked for. + logger.warning("AutoEP: the DeepEP backend does not support folded tensor parallelism; " + "using the default all-to-all for this layer") + self.comm_backend = NCCL_BACKEND self.ep_group_name = folding_group_handles.ep_group_name self.ep_group = folding_group_handles.ep_group self.tp_group = folding_group_handles.tp_group @@ -593,10 +601,13 @@ def _deepep_route(self, tokens: torch.Tensor, ro: "RouterOutput") -> torch.Tenso expansion and the reduction around the collectives, not just the collectives, and returns [T, H] ready for the shared tail of forward. """ - if self._deepep_exchange is None: - # Built on first use rather than at construction, because the - # buffer is sized from the token count and that is not known until - # a batch arrives. + # Built on first use because the buffer is sized from the token count, + # 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: + if self._deepep_exchange is not None: + self._deepep_exchange.destroy() self._deepep_exchange = DeepEPExchange( ep_group=self.ep_group, num_experts=self.num_experts, @@ -605,7 +616,8 @@ def _deepep_route(self, tokens: torch.Tensor, ro: "RouterOutput") -> torch.Tenso num_max_tokens_per_rank=tokens.shape[0], ) - received, exchange = deepep_dispatch(self._deepep_exchange, tokens, ro.selected_experts, ro.top_scores) + received, recv_weights, exchange = deepep_dispatch(self._deepep_exchange, tokens, ro.selected_experts, + ro.top_scores) handle = exchange.last_handle # Two quantities have to line up for the grouped GEMM, and they come @@ -626,7 +638,6 @@ def _deepep_route(self, tokens: torch.Tensor, ro: "RouterOutput") -> torch.Tenso # the reduction and is handed to combine instead. In expand mode each # row is one token-expert slot, so a weight per row broadcasts across # the hidden dimension. - recv_weights = exchange.last_recv_weights if self.score_apply == "pre" and recv_weights is not None: scale = recv_weights[:received.shape[0]].reshape(-1, 1) received = (received.to(torch.float32) * scale).to(received.dtype) @@ -666,6 +677,9 @@ def forward( expert_indices_sorted = ro.selected_experts.reshape(-1).index_select(0, token_indices_sorted) folded_tp = self.folding_group_handles is not None and self.folding_group_handles.spec.tp_size > 1 + # Set only where DeepEP's combine actually produced the output, since + # that decides whether the reduction below has already happened. + deepep_combined = False restore_ctx = None if folded_tp: from deepspeed.moe.ep_tp_dispatch import ( @@ -745,6 +759,7 @@ def forward( if self.comm_backend == DEEPEP_BACKEND: expert_output = self._deepep_route(x, ro) + deepep_combined = True else: routed_input = _AllToAllV.apply(self.ep_group, routed_input, plan.input_splits, plan.output_splits) @@ -761,8 +776,11 @@ def forward( tp_group=self.tp_group, validate_coverage=self.validate_folding_routing).reshape(bsz, seqlen, hdim) self._last_folding_dispatch_counters = dispatch_counters(restore_ctx) - elif self.comm_backend == DEEPEP_BACKEND: - # Already reduced over top-k and back in token order. + elif deepep_combined: + # DeepEP's combine already reduced over top-k and restored token + # order. This is keyed on the route having run rather than on the + # backend being selected: with ep_size == 1 the local path runs + # instead and still has one row per assignment to reduce. output = expert_output.reshape(bsz, seqlen, hdim) else: output = combine_from_routed( diff --git a/docs/code-docs/source/autoep.rst b/docs/code-docs/source/autoep.rst index b7b7fa293d82..29f05edd978d 100644 --- a/docs/code-docs/source/autoep.rst +++ b/docs/code-docs/source/autoep.rst @@ -84,6 +84,37 @@ Weights-only/module-only Universal Checkpoint loads use the converted 4. Expert parameters are marked for expert-data-parallel gradient reduction; router and shared-expert parameters use standard data-parallel reduction. +**Communication backend (optional):** + +The expert AllToAll can be carried by `DeepEP `__ +instead of NCCL. This is opt-in and off by default; jobs that set nothing keep +the NCCL path unchanged. + +.. code-block:: bash + + export DEEPSPEED_AUTOEP_COMM_BACKEND=deepep # default: nccl + export DEEPSPEED_AUTOEP_COMM_SMS=12 # communication SM budget + +On 16 H100s across two nodes, replaying routing captured from real training, +DeepEP reduced payload AllToAll time from roughly 100 ms to 48 ms per step, and +whole training steps from 454 ms to 371 ms. The advantage grows with routing +imbalance: at the most skewed step measured, the NCCL path degraded to 116 ms +while DeepEP stayed flat. + +``DEEPSPEED_AUTOEP_COMM_SMS`` matters because communication competes with the +expert GEMM for SMs. The default of 12 was chosen by measuring whole steps: a +smaller budget slows the collective itself, and a larger one slows everything +else. Requirements and limits: + +- The ``deep_ep`` package must be installed. It is imported only when this + backend is selected, so installations without it are unaffected. +- DeepEP v2 requires NCCL 2.30.4 or newer, built with GIN support. Below that + version the transport is unavailable regardless of the network. +- DeepEP v1 (the legacy ``Buffer`` API, using NVSHMEM and IBGDA) is not + supported. +- The backend is disabled, with a warning, on layers using folded tensor + parallelism (``tp_size > 1``); those layers fall back to NCCL. + **Constraints:** - ``autoep_size`` must divide ``num_experts`` for all detected MoE layers. diff --git a/tests/unit/module_inject/test_auto_ep_comm.py b/tests/unit/module_inject/test_auto_ep_comm.py index 6d4904e675a7..3d273062ba8d 100644 --- a/tests/unit/module_inject/test_auto_ep_comm.py +++ b/tests/unit/module_inject/test_auto_ep_comm.py @@ -1,7 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # DeepSpeed Team +import ast +import inspect import os +import textwrap import sys import unittest @@ -9,7 +12,8 @@ from unittest import mock from deepspeed.module_inject.auto_ep_comm import (AVAILABLE_BACKENDS, DEEPEP_BACKEND, NCCL_BACKEND, _conform_rows, - _import_deep_ep, configured_backend, configured_num_sms) + _DeepEPCombine, _DeepEPDispatch, _import_deep_ep, configured_backend, + configured_num_sms) class TestAutoEPCommBackendSelection(unittest.TestCase): @@ -114,5 +118,47 @@ def test_matching_shape_is_passed_through_untouched(self): self.assertIs(_conform_rows(grad, (7, 4)), grad) +class TestAutogradSignatures(unittest.TestCase): + """Both directions must return a gradient for every differentiable input. + + A missing router-weight gradient does not fail loudly: training runs, the + loss falls, and the gate simply never learns. These check the arity that + carries it rather than leaving it to a live run to reveal. + """ + + @staticmethod + def gradient_count(function) -> int: + """How many values the backward returns, parsed rather than counted. + + Counting commas in the source would also count the ones inside calls + like ``_conform_rows(grad, shape)``. + """ + tree = ast.parse(textwrap.dedent(inspect.getsource(function))) + returns = [node for node in ast.walk(tree) if isinstance(node, ast.Return)] + assert returns, "backward has no return statement" + value = returns[-1].value + return len(value.elts) if isinstance(value, ast.Tuple) else 1 + + def test_dispatch_backward_returns_a_gradient_per_input(self): + # ctx is not an input autograd returns a gradient for. + inputs = len(inspect.signature(_DeepEPDispatch.forward).parameters) - 1 + + self.assertEqual(inputs, 4) + self.assertEqual(self.gradient_count(_DeepEPDispatch.backward), inputs) + + def test_combine_backward_returns_a_gradient_per_input(self): + inputs = len(inspect.signature(_DeepEPCombine.forward).parameters) - 1 + + self.assertEqual(inputs, 4) + self.assertEqual(self.gradient_count(_DeepEPCombine.backward), inputs) + + def test_dispatch_forward_returns_weights_so_they_stay_differentiable(self): + # The received weights have to leave the custom function as an output; + # reading them off the exchange afterwards puts them outside the graph. + source = inspect.getsource(_DeepEPDispatch.forward) + + self.assertIn("return received, recv_weights", source) + + if __name__ == "__main__": unittest.main() From 12feb411be0a90985599454e78057f11ead47bb1 Mon Sep 17 00:00:00 2001 From: yh0903 Date: Tue, 4 Aug 2026 11:19:26 -0700 Subject: [PATCH 3/7] Drop the exchange field the weights no longer travel through 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 --- deepspeed/module_inject/auto_ep_comm.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/deepspeed/module_inject/auto_ep_comm.py b/deepspeed/module_inject/auto_ep_comm.py index 5b2d03a05c8c..4016a9712892 100644 --- a/deepspeed/module_inject/auto_ep_comm.py +++ b/deepspeed/module_inject/auto_ep_comm.py @@ -173,9 +173,6 @@ def __init__(self, ep_group, num_experts: int, top_k: int, hidden_size: int, num # The handle the last dispatch produced. Combine and both backward # passes replay against it, so it has to outlive the dispatch call. self.last_handle = None - # The routing weights that arrived with the last dispatch, kept so the - # layer can decide whether to apply them before or after the experts. - self.last_recv_weights = None def dispatch(self, tokens: torch.Tensor, topk_idx: torch.Tensor, topk_weights: torch.Tensor): """Send tokens to their experts, returning rows, weights and handle. @@ -203,7 +200,6 @@ def dispatch(self, tokens: torch.Tensor, topk_idx: torch.Tensor, topk_weights: t # The returned event only holds anything when the call was made with # async_with_compute_stream; a synchronous result is already usable. self.last_handle = handle - self.last_recv_weights = recv_weights return recv_x, recv_weights, handle def dispatch_with_handle(self, tokens: torch.Tensor, handle): From d2071c072817191079583141e1ee141ee8564497 Mon Sep 17 00:00:00 2001 From: yh0903 Date: Wed, 5 Aug 2026 03:16:53 -0700 Subject: [PATCH 4/7] Agree DeepEP buffer capacity across the group and trim weights with rows 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 --- deepspeed/module_inject/auto_ep_comm.py | 6 +- deepspeed/module_inject/auto_ep_layer.py | 24 +++++- tests/unit/module_inject/test_auto_ep_comm.py | 73 +++++++++++++++++++ 3 files changed, 97 insertions(+), 6 deletions(-) diff --git a/deepspeed/module_inject/auto_ep_comm.py b/deepspeed/module_inject/auto_ep_comm.py index 4016a9712892..7e498ed1a9b8 100644 --- a/deepspeed/module_inject/auto_ep_comm.py +++ b/deepspeed/module_inject/auto_ep_comm.py @@ -255,7 +255,7 @@ def _conform_rows(tensor: torch.Tensor, shape) -> torch.Tensor: return tensor if tensor.shape[0] > rows: return tensor[:rows] - extended = tensor.new_zeros((rows, tensor.shape[1])) + extended = tensor.new_zeros((rows, ) + tuple(tensor.shape[1:])) extended[:tensor.shape[0]] = tensor return extended @@ -286,7 +286,7 @@ def backward(ctx, grad_received, grad_recv_weights): ) conformed_weights = None if grad_weights is not None and ctx.weights_shape is not None: - conformed_weights = grad_weights[:ctx.weights_shape[0]].reshape(ctx.weights_shape) + conformed_weights = _conform_rows(grad_weights, ctx.weights_shape).reshape(ctx.weights_shape) return None, _conform_rows(grad_tokens, ctx.tokens_shape), None, conformed_weights @@ -309,7 +309,7 @@ def backward(ctx, grad_combined): grad_rows, grad_weights = ctx.exchange.dispatch_with_handle(grad_combined.contiguous(), ctx.handle) conformed_weights = None if grad_weights is not None and ctx.weights_shape is not None: - conformed_weights = grad_weights[:ctx.weights_shape[0]] + conformed_weights = _conform_rows(grad_weights, ctx.weights_shape) return None, _conform_rows(grad_rows, ctx.rows_shape), None, conformed_weights diff --git a/deepspeed/module_inject/auto_ep_layer.py b/deepspeed/module_inject/auto_ep_layer.py index 3dbfcc636129..3c594528fd24 100644 --- a/deepspeed/module_inject/auto_ep_layer.py +++ b/deepspeed/module_inject/auto_ep_layer.py @@ -605,7 +605,19 @@ def _deepep_route(self, tokens: torch.Tensor, ro: "RouterOutput") -> torch.Tenso # 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: + # + # Capacity is agreed across the group before it is used. Ranks hold + # different local token counts, so sizing from the local count has some + # ranks constructing a buffer while others do not; construction is + # collective, and the ranks that entered it then wait for peers that + # never arrive until the connection drops. Reducing first also makes + # the rebuild decision unanimous for free, since every rank compares + # the same global count against the same current capacity. + capacity = torch.tensor([tokens.shape[0]], dtype=torch.int64, device=tokens.device) + dist.all_reduce(capacity, op=dist.ReduceOp.MAX, group=self.ep_group) + num_max_tokens_per_rank = int(capacity.item()) + + if self._deepep_exchange is None or num_max_tokens_per_rank > self._deepep_exchange.num_max_tokens_per_rank: if self._deepep_exchange is not None: self._deepep_exchange.destroy() self._deepep_exchange = DeepEPExchange( @@ -613,7 +625,7 @@ def _deepep_route(self, tokens: torch.Tensor, ro: "RouterOutput") -> torch.Tenso num_experts=self.num_experts, top_k=self.top_k, hidden_size=self.hidden_size, - num_max_tokens_per_rank=tokens.shape[0], + num_max_tokens_per_rank=num_max_tokens_per_rank, ) received, recv_weights, exchange = deepep_dispatch(self._deepep_exchange, tokens, ro.selected_experts, @@ -632,6 +644,12 @@ def _deepep_route(self, tokens: torch.Tensor, ro: "RouterOutput") -> torch.Tenso raise RuntimeError(f"DeepEP returned {counts.numel()} expert counts, but this rank owns " f"{self.num_local_experts} experts") + # The weights arrive in a buffer sized for the worst case, like the + # rows did, so they are trimmed to the same prefix. Combine reduces + # rows and weights together and rejects a pair whose lengths disagree. + if recv_weights is not None: + recv_weights = recv_weights[:received.shape[0]] + # The router weights must be applied exactly once. In "pre" mode the # NCCL path scales tokens before the experts see them, so the arrived # rows are scaled the same way; in "post" mode the scaling is part of @@ -639,7 +657,7 @@ def _deepep_route(self, tokens: torch.Tensor, ro: "RouterOutput") -> torch.Tenso # row is one token-expert slot, so a weight per row broadcasts across # the hidden dimension. if self.score_apply == "pre" and recv_weights is not None: - scale = recv_weights[:received.shape[0]].reshape(-1, 1) + scale = recv_weights.reshape(-1, 1) received = (received.to(torch.float32) * scale).to(received.dtype) combine_weights = None else: diff --git a/tests/unit/module_inject/test_auto_ep_comm.py b/tests/unit/module_inject/test_auto_ep_comm.py index 3d273062ba8d..9cd80e6854bb 100644 --- a/tests/unit/module_inject/test_auto_ep_comm.py +++ b/tests/unit/module_inject/test_auto_ep_comm.py @@ -117,6 +117,79 @@ def test_matching_shape_is_passed_through_untouched(self): self.assertIs(_conform_rows(grad, (7, 4)), grad) + def test_conforms_a_one_dimensional_weight_buffer(self): + # Routing weights arrive one per row rather than one per hidden unit, + # so conforming has to work without a trailing dimension. + grad = torch.ones(9) + + self.assertEqual(tuple(_conform_rows(grad, (4, )).shape), (4, )) + self.assertEqual(tuple(_conform_rows(grad, (12, )).shape), (12, )) + + +class TestRoutedRowAgreement(unittest.TestCase): + """Rows and their weights must describe the same arrivals. + + DeepEP hands back whole worst-case buffers, and the layer trims the rows to + the arrivals the prefix sum reports. Leaving the weights at full length + makes combine reduce a row count and a weight count that disagree, which is + silent in "pre" mode and fatal in "post" mode. + """ + + BUFFER_ROWS = 32 + ARRIVED_ROWS = 12 + HIDDEN = 8 + LOCAL_EXPERTS = 4 + + def route(self, score_apply): + """Drive _deepep_route with a stub exchange, returning combine's inputs.""" + from deepspeed.module_inject import auto_ep_layer + + prefix = torch.tensor([3, 6, 9, self.ARRIVED_ROWS], dtype=torch.int64) + handle = mock.Mock(psum_num_recv_tokens_per_expert=prefix) + exchange = mock.Mock(last_handle=handle, num_max_tokens_per_rank=self.BUFFER_ROWS) + + buffer_rows = torch.ones((self.BUFFER_ROWS, self.HIDDEN)) + buffer_weights = torch.ones(self.BUFFER_ROWS) + seen = {} + + def fake_dispatch(_exchange, *_args): + return buffer_rows, buffer_weights, exchange + + def fake_combine(_exchange, rows, _handle, topk_weights=None): + seen["rows"] = rows + seen["weights"] = topk_weights + return rows + + layer = mock.Mock( + _deepep_exchange=exchange, + num_local_experts=self.LOCAL_EXPERTS, + score_apply=score_apply, + experts=lambda rows, counts: rows, + ) + + with mock.patch.object(auto_ep_layer, "deepep_dispatch", fake_dispatch), \ + mock.patch.object(auto_ep_layer, "deepep_combine", fake_combine), \ + mock.patch.object(auto_ep_layer.dist, "all_reduce", lambda *a, **k: None): + router_output = auto_ep_layer.RouterOutput( + top_scores=torch.ones((4, 2)), + selected_experts=torch.zeros((4, 2), dtype=torch.long), + num_tokens_per_expert=torch.zeros(self.LOCAL_EXPERTS, dtype=torch.long), + ) + auto_ep_layer.AutoEPMoELayer._deepep_route(layer, torch.ones((4, self.HIDDEN)), router_output) + return seen + + def test_post_mode_hands_combine_one_weight_per_row(self): + seen = self.route("post") + + self.assertEqual(seen["rows"].shape[0], self.ARRIVED_ROWS) + self.assertEqual(seen["weights"].shape[0], self.ARRIVED_ROWS) + + def test_pre_mode_scales_only_the_rows_that_arrived(self): + seen = self.route("pre") + + self.assertEqual(seen["rows"].shape[0], self.ARRIVED_ROWS) + self.assertIsNone(seen["weights"], "pre mode applies the weights before the experts, not during combine") + class TestAutogradSignatures(unittest.TestCase): """Both directions must return a gradient for every differentiable input. From e6fbaf942d4c1ce94cb107ce900a39cc22b5a501 Mon Sep 17 00:00:00 2001 From: yh0903 Date: Wed, 5 Aug 2026 10:27:08 -0700 Subject: [PATCH 5/7] Address maintainer review on the DeepEP backend 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 --- deepspeed/module_inject/auto_ep_comm.py | 172 ++++++++-------- deepspeed/module_inject/auto_ep_config.py | 17 ++ deepspeed/module_inject/auto_ep_layer.py | 114 +++++----- .../module_inject/auto_ep_presets/base.py | 3 + deepspeed/runtime/engine.py | 5 + docs/code-docs/source/autoep.rst | 44 ++-- tests/unit/module_inject/test_auto_ep_comm.py | 194 +++++++++++++----- 7 files changed, 354 insertions(+), 195 deletions(-) diff --git a/deepspeed/module_inject/auto_ep_comm.py b/deepspeed/module_inject/auto_ep_comm.py index 7e498ed1a9b8..86645c4d5928 100644 --- a/deepspeed/module_inject/auto_ep_comm.py +++ b/deepspeed/module_inject/auto_ep_comm.py @@ -4,15 +4,16 @@ The dispatch and combine collectives are the largest single cost in an AutoEP step, and a measured replay of real routing on 16 H100s across two nodes put -NCCL at 195.8 ms of payload all-to-all per step against DeepEP's 86.3 ms. This -module is what lets that be switched without the MoE layer knowing which -transport it is using. +the collective path at 195.8 ms of payload all-to-all per step against DeepEP's +86.3 ms. This module is what lets that be switched without the MoE layer +knowing which transport it is using. -Selection is by environment variable and defaults to the NCCL path, so a job -that sets nothing behaves exactly as before: +Selection lives in the ``expert_parallel`` section of the DeepSpeed config and +defaults to the collective path, so a job that sets nothing behaves exactly as +before:: - DEEPSPEED_AUTOEP_COMM_BACKEND=nccl (default) - DEEPSPEED_AUTOEP_COMM_BACKEND=deepep + "expert_parallel": {"comm_backend": "comm"} # default + "expert_parallel": {"comm_backend": "deepep"} A backend is asked for once per layer and reused, because DeepEP's buffers are sized at construction and are expensive to rebuild. @@ -20,22 +21,21 @@ from __future__ import annotations -import os - import torch from deepspeed.utils import logger -_BACKEND_ENV = "DEEPSPEED_AUTOEP_COMM_BACKEND" -_NUM_SMS_ENV = "DEEPSPEED_AUTOEP_COMM_SMS" -NCCL_BACKEND = "nccl" +# Names the transport, not the library behind it: the default path goes through +# deepspeed.comm, which is NCCL on CUDA but not on every accelerator. +COMM_BACKEND = "comm" # Names the library, not its version: v2 is the only path implemented, and # nothing about the name would have to change if that ever grew. DEEPEP_BACKEND = "deepep" -AVAILABLE_BACKENDS = (NCCL_BACKEND, DEEPEP_BACKEND) +AVAILABLE_BACKENDS = (COMM_BACKEND, DEEPEP_BACKEND) # GIN, and so DeepEP, does not exist in any form below this NCCL version. NCCL_GIN_MIN_VERSION = (2, 30, 4) -# Chosen by sweeping whole training steps rather than the collective alone. +# The config's comm_num_sm default. Chosen by sweeping whole training steps +# rather than the collective alone. # On 16 H100s across two nodes the step took 340, 311, 353, 360 and 391 ms at 8, # 12, 16, 24 and 32 SMs. Twelve is the point where the collective is already as # fast as it gets while the expert GEMM still has the SMs it needs: at 8 the @@ -44,45 +44,44 @@ DEFAULT_COMM_SMS = 12 -def _qps_for_sms(num_sms: int) -> int: +def _qps_for_sms(num_sms: int, qp_margin: int) -> int: """Queue pairs to reserve for a given SM count. - One per SM plus a small margin for the control path. This is deliberately - smaller than DeepEP's automatic choice, which assumes it is the only thing - on the fabric. + One per SM plus a margin for the control path. This is deliberately smaller + than DeepEP's automatic choice, which assumes it is the only thing on the + fabric: in a training step ZeRO and the data-parallel groups have already + taken their share, and asking for DeepEP's default exhausts them. """ - return num_sms + 4 + return num_sms + qp_margin -def configured_backend() -> str: - """The transport this process should use for expert all-to-all. +# Every buffer built in this process, in construction order. DeepEP buffers are +# constructed with explicitly_destroy, so nothing reclaims them on its own, and +# destroying them is collective: every rank has to do it in the same order. +_LIVE_EXCHANGES: list["DeepEPExchange"] = [] - An unset variable selects NCCL, so a job that opts into nothing keeps the - behaviour it had before this existed. - """ - raw = os.environ.get(_BACKEND_ENV) - if raw is None or not raw.strip(): - return NCCL_BACKEND - name = raw.strip().lower() - if name not in AVAILABLE_BACKENDS: - raise ValueError(f"{_BACKEND_ENV}={raw!r} is not one of {AVAILABLE_BACKENDS}") - return name +# DeepEP's dispatch kernel handles bfloat16 and fp8, not fp16. A half-precision +# run would otherwise reach an assertion inside the kernel. +SUPPORTED_DTYPES = (torch.bfloat16, torch.float32) + + +def assert_dtype_supported(dtype: torch.dtype) -> None: + """Reject dtypes DeepEP's kernels cannot dispatch.""" + if dtype not in SUPPORTED_DTYPES: + raise TypeError(f'comm_backend="{DEEPEP_BACKEND}" does not support {dtype}: DeepEP\'s dispatch kernel ' + 'handles bfloat16, not fp16. Train in bfloat16, or set comm_backend="comm" to use the ' + "default all-to-all, which has no such restriction.") -def configured_num_sms() -> int: - """SM budget for communication; 0 lets the backend decide. +def destroy_all_exchanges() -> None: + """Release every DeepEP buffer this process built. - Worth setting because the collective competes with the expert GEMM for - SMs: measurements showed the GEMM running 1.21-1.25x slower whenever a - collective was in flight. + Collective, and ordered by construction, so every rank tears the same + buffers down in the same order. Worth calling at the end of training: the + buffers ask DeepEP not to reclaim them, so nothing else will. """ - raw = os.environ.get(_NUM_SMS_ENV) - if raw is None or not raw.strip(): - return 0 - try: - return int(raw) - except ValueError as error: - raise ValueError(f"{_NUM_SMS_ENV}={raw!r} is not an integer") from error + for exchange in list(_LIVE_EXCHANGES): + exchange.destroy() def _import_deep_ep(): @@ -98,10 +97,10 @@ def _import_deep_ep(): import deep_ep except ImportError as error: raise ImportError( - f"{_BACKEND_ENV}={DEEPEP_BACKEND} requires the deep_ep package, which is not installed. It also " + f'comm_backend="{DEEPEP_BACKEND}" requires the deep_ep package, which is not installed. It also ' "requires NCCL 2.30.4 or newer built with GIN support: the transport is unavailable below that " - "version regardless of the network. Unset " - f"{_BACKEND_ENV} to use the default NCCL all-to-all, which has no such requirement.") from error + 'version regardless of the network. Set comm_backend="comm" to use the default all-to-all, which ' + "has no such requirement.") from error nccl_version = _nccl_version() if nccl_version is not None and nccl_version < NCCL_GIN_MIN_VERSION: @@ -115,7 +114,7 @@ def _import_deep_ep(): logger.warning( f"torch reports NCCL {installed}, older than the {minimum} that GIN requires. DeepEP links its own " "NCCL, so this is only a problem if it also resolves to the older one; a failure inside buffer " - f"construction is the symptom. Unset {_BACKEND_ENV} to fall back to the default all-to-all.") + 'construction is the symptom. Set comm_backend="comm" to fall back to the default all-to-all.') return deep_ep @@ -142,7 +141,14 @@ class DeepEPExchange: survive from forward to backward. """ - def __init__(self, ep_group, num_experts: int, top_k: int, hidden_size: int, num_max_tokens_per_rank: int): + def __init__(self, + ep_group, + num_experts: int, + top_k: int, + hidden_size: int, + num_max_tokens_per_rank: int, + num_sms: int = DEFAULT_COMM_SMS, + qp_margin: int = 4): deep_ep = _import_deep_ep() self.deep_ep = deep_ep @@ -151,14 +157,13 @@ def __init__(self, ep_group, num_experts: int, top_k: int, hidden_size: int, num # nothing else but fails in a training step where ZeRO and the # data-parallel groups have already taken their share. Asking for only # what the chosen SM count needs keeps the request proportionate. - num_sms = configured_num_sms() or DEFAULT_COMM_SMS self.buffer = deep_ep.ElasticBuffer( ep_group, num_max_tokens_per_rank=num_max_tokens_per_rank, hidden=hidden_size, num_topk=top_k, use_fp8_dispatch=False, - num_allocated_qps=_qps_for_sms(num_sms), + num_allocated_qps=_qps_for_sms(num_sms, qp_margin), # Required once the EP group spans nodes: it splits the ranks into # an NVLink domain and an RDMA domain rather than assuming a single # flat NVLink domain. @@ -166,13 +171,17 @@ def __init__(self, ep_group, num_experts: int, top_k: int, hidden_size: int, num explicitly_destroy=True, ) self.num_sms = num_sms - self.num_qps = self.buffer.get_theoretical_num_qps(self.num_sms) self.num_experts = num_experts # Recorded so the layer can tell when a later batch outgrows it. self.num_max_tokens_per_rank = num_max_tokens_per_rank # The handle the last dispatch produced. Combine and both backward # passes replay against it, so it has to outlive the dispatch call. self.last_handle = None + self.destroyed = False + # Constructed with explicitly_destroy, so nothing reclaims this buffer + # on its own. Registering it means a process that never calls the + # layer's teardown can still release every buffer in one call. + _LIVE_EXCHANGES.append(self) def dispatch(self, tokens: torch.Tensor, topk_idx: torch.Tensor, topk_weights: torch.Tensor): """Send tokens to their experts, returning rows, weights and handle. @@ -195,32 +204,32 @@ def dispatch(self, tokens: torch.Tensor, topk_idx: torch.Tensor, topk_weights: t # offsets have to describe the rows that are actually there. expert_alignment=1, num_sms=self.num_sms, - num_qps=self.num_qps, ) # The returned event only holds anything when the call was made with # async_with_compute_stream; a synchronous result is already usable. self.last_handle = handle return recv_x, recv_weights, handle - def dispatch_with_handle(self, tokens: torch.Tensor, handle): - """Replay a dispatch against a cached handle, keeping both outputs. + def dispatch_with_handle(self, tokens: torch.Tensor, handle) -> torch.Tensor: + """Replay a dispatch against a cached handle. - Used as the backward of a combine, where the weights matter as much as - the rows: their gradient is what reaches the router gate. + Used as the backward of a combine, which scatters the combined + gradient back to the rows that contributed to it. """ - recv_x, _, recv_weights, _, _ = self.buffer.dispatch( + recv_x, _, _, _, _ = self.buffer.dispatch( tokens, handle=handle, num_sms=self.num_sms, - num_qps=self.num_qps, ) - return recv_x, recv_weights + return recv_x def combine_with_weight_grad(self, rows: torch.Tensor, handle, weight_grads=None): - """Combine that also returns the reduced weight gradient. + """Combine that also reduces the routing-weight gradient. - Used as the backward of a dispatch. The weight gradient travels the - same path as the row gradient, so it is reduced by the same call. + Used as the backward of a dispatch. Dispatch replicates a token's + routing weight to every rank that expert-owns it, so the adjoint is a + sum over those copies, which is exactly what combine does to the + weights it carries. """ combined, combined_weights, _ = self.buffer.combine(rows, handle=handle, @@ -228,18 +237,26 @@ def combine_with_weight_grad(self, rows: torch.Tensor, handle, weight_grads=None num_sms=self.num_sms) return combined, combined_weights - def combine(self, rows: torch.Tensor, handle, topk_weights=None) -> torch.Tensor: - """Reduce expert outputs back to their tokens. + def combine(self, rows: torch.Tensor, handle) -> torch.Tensor: + """Reduce expert outputs back to the tokens they came from. - Passing ``topk_weights`` makes DeepEP weight each expert's output as it - reduces, which is the same weighted sum the NCCL path performs - separately after its combine. + Deliberately does not pass ``topk_weights``. DeepEP's combine does not + multiply the rows by those weights; it transports and reduces them + alongside, returning them separately. Handing the routing weights here + would therefore drop them from the result, so the layer applies them to + the rows itself. """ - combined, _, _ = self.buffer.combine(rows, handle=handle, topk_weights=topk_weights, num_sms=self.num_sms) + combined, _, _ = self.buffer.combine(rows, handle=handle, num_sms=self.num_sms) return combined def destroy(self) -> None: + """Release the buffer. Collective, so every rank must call it.""" + if self.destroyed: + return + self.destroyed = True self.buffer.destroy() + if self in _LIVE_EXCHANGES: + _LIVE_EXCHANGES.remove(self) def _conform_rows(tensor: torch.Tensor, shape) -> torch.Tensor: @@ -294,23 +311,18 @@ class _DeepEPCombine(torch.autograd.Function): """Combine whose backward is the matching dispatch, on the same handle.""" @staticmethod - def forward(ctx, exchange: DeepEPExchange, rows: torch.Tensor, handle, topk_weights): + def forward(ctx, exchange: DeepEPExchange, rows: torch.Tensor, handle): ctx.exchange = exchange ctx.handle = handle - # The forward input was trimmed to the rows that actually arrived, - # while the backward dispatch hands back a whole buffer. Autograd + # The backward dispatch hands back a whole buffer, while autograd # requires the gradient to match the input it is the gradient of. ctx.rows_shape = rows.shape - ctx.weights_shape = None if topk_weights is None else topk_weights.shape - return exchange.combine(rows, handle, topk_weights) + return exchange.combine(rows, handle) @staticmethod def backward(ctx, grad_combined): - grad_rows, grad_weights = ctx.exchange.dispatch_with_handle(grad_combined.contiguous(), ctx.handle) - conformed_weights = None - if grad_weights is not None and ctx.weights_shape is not None: - conformed_weights = _conform_rows(grad_weights, ctx.weights_shape) - return None, _conform_rows(grad_rows, ctx.rows_shape), None, conformed_weights + grad_rows = ctx.exchange.dispatch_with_handle(grad_combined.contiguous(), ctx.handle) + return None, _conform_rows(grad_rows, ctx.rows_shape), None def deepep_dispatch(exchange: DeepEPExchange, tokens: torch.Tensor, topk_idx: torch.Tensor, @@ -320,5 +332,5 @@ def deepep_dispatch(exchange: DeepEPExchange, tokens: torch.Tensor, topk_idx: to return received, recv_weights, exchange -def deepep_combine(exchange: DeepEPExchange, rows: torch.Tensor, handle, topk_weights=None) -> torch.Tensor: - return _DeepEPCombine.apply(exchange, rows, handle, topk_weights) +def deepep_combine(exchange: DeepEPExchange, rows: torch.Tensor, handle) -> torch.Tensor: + return _DeepEPCombine.apply(exchange, rows, handle) diff --git a/deepspeed/module_inject/auto_ep_config.py b/deepspeed/module_inject/auto_ep_config.py index 2841b5317f8c..816ffa61931d 100644 --- a/deepspeed/module_inject/auto_ep_config.py +++ b/deepspeed/module_inject/auto_ep_config.py @@ -58,6 +58,9 @@ def parse_autoep_config(param_dict: dict) -> AutoEPConfig: config.route_scale = param_dict.get("route_scale", 1.0) config.score_apply = param_dict.get("score_apply", "auto") config.combine_impl = param_dict.get("combine_impl", "auto") + config.comm_backend = param_dict.get("comm_backend", "comm") + config.comm_num_sm = param_dict.get("comm_num_sm", 12) + config.comm_qp_margin = param_dict.get("comm_qp_margin", 4) config.num_expert_groups = param_dict.get("num_expert_groups", None) config.num_limited_groups = param_dict.get("num_limited_groups", None) config.score_func = param_dict.get("score_func", "auto") @@ -157,6 +160,20 @@ def validate_autoep_config( raise ValueError(f"combine_impl must be one of {valid_combine_impl}, " f"got '{config.combine_impl}'") + # Validate comm_backend + valid_comm_backend = ("comm", "deepep") + if config.comm_backend not in valid_comm_backend: + raise ValueError(f"comm_backend must be one of {valid_comm_backend}, " + f"got '{config.comm_backend}'") + + # A zero budget would hand the whole GPU to the collective, and a negative + # one is meaningless; both are worth rejecting where the value is written + # rather than inside a buffer constructor. + if not isinstance(config.comm_num_sm, int) or config.comm_num_sm < 1: + raise ValueError(f"comm_num_sm must be a positive integer, got {config.comm_num_sm!r}") + if not isinstance(config.comm_qp_margin, int) or config.comm_qp_margin < 0: + raise ValueError(f"comm_qp_margin must be a non-negative integer, got {config.comm_qp_margin!r}") + # Validate score_func valid_score_func = ("auto", "softmax", "sigmoid") if config.score_func not in valid_score_func: diff --git a/deepspeed/module_inject/auto_ep_layer.py b/deepspeed/module_inject/auto_ep_layer.py index 3c594528fd24..bf0c64226597 100644 --- a/deepspeed/module_inject/auto_ep_layer.py +++ b/deepspeed/module_inject/auto_ep_layer.py @@ -22,7 +22,7 @@ from deepspeed.module_inject.auto_ep_config import AutoEPConfig, MoELayerSpec, resolve_autoep_config_defaults from deepspeed.module_inject.auto_ep_folding import mark_autoep_folding_router_parameter from deepspeed.utils import logger -from deepspeed.module_inject.auto_ep_comm import (DEEPEP_BACKEND, NCCL_BACKEND, DeepEPExchange, configured_backend, +from deepspeed.module_inject.auto_ep_comm import (DEEPEP_BACKEND, DeepEPExchange, assert_dtype_supported, deepep_combine, deepep_dispatch) from deepspeed.moe.ep_router import TokenChoiceTopKRouter from deepspeed.moe.ep_count import count_tokens_per_expert @@ -523,14 +523,16 @@ def __init__( # Router-logit cache self._cached_router_logits = None # Resolved once per layer: a DeepEP exchange sizes its buffer at - # construction, so it is built on first use and kept. A bad value must - # not break the default path, which nobody opted into. - try: - self.comm_backend = configured_backend() - except ValueError as error: - logger.warning(f"AutoEP: falling back to the default all-to-all: {error}") - self.comm_backend = NCCL_BACKEND + # construction, so it is built on first use and kept. + self.comm_backend = config.comm_backend + self.comm_num_sm = config.comm_num_sm + self.comm_qp_margin = config.comm_qp_margin self._deepep_exchange = None + # Buffers outgrown by a larger batch. They stay alive rather than being + # released, because a backward still pending from an earlier forward + # replays against the buffer that forward used. Capacity is rounded up + # so this list stays short. + self._retired_exchanges = [] self._register_logit_hook() def _register_logit_hook(self): @@ -564,11 +566,13 @@ def set_deepspeed_parallelism( if self.comm_backend == DEEPEP_BACKEND and folding_group_handles.spec.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. Saying so beats - # silently running something other than what was asked for. - logger.warning("AutoEP: the DeepEP backend does not support folded tensor parallelism; " - "using the default all-to-all for this layer") - self.comm_backend = NCCL_BACKEND + # returns token-major rows cannot satisfy. Refusing beats + # falling back, which leaves the job reporting a backend it is + # not running and the speedup unexplained. + raise ValueError( + f'comm_backend="{DEEPEP_BACKEND}" does not support folded tensor parallelism ' + f"(expert_tensor_parallel_size={folding_group_handles.spec.tp_size}). Set " + 'expert_tensor_parallel_size to 1, or comm_backend to "comm".') self.ep_group_name = folding_group_handles.ep_group_name self.ep_group = folding_group_handles.ep_group self.tp_group = folding_group_handles.tp_group @@ -597,74 +601,84 @@ def _deepep_route(self, tokens: torch.Tensor, ro: "RouterOutput") -> torch.Tenso ``tokens`` is [T, H] before top-k expansion. DeepEP replicates each token to the ranks that need it rather than being handed one row per selected expert, groups arrivals by expert for the grouped GEMM, and - reduces the weighted sum on the way back. It therefore replaces the - expansion and the reduction around the collectives, not just the - collectives, and returns [T, H] ready for the shared tail of forward. + sums them back. It therefore replaces the expansion and the reduction + around the collectives, not just the collectives, and returns [T, H] + ready for the shared tail of forward. """ + assert_dtype_supported(tokens.dtype) + # Built on first use because the buffer is sized from the token count, # 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. # # Capacity is agreed across the group before it is used. Ranks hold - # different local token counts, so sizing from the local count has some - # ranks constructing a buffer while others do not; construction is - # collective, and the ranks that entered it then wait for peers that - # never arrive until the connection drops. Reducing first also makes - # the rebuild decision unanimous for free, since every rank compares - # the same global count against the same current capacity. + # different local token counts, and DeepEP requires the same + # num_max_tokens_per_rank everywhere; sizing from the local count also + # has some ranks entering construction while others do not, and + # construction is collective, so those that entered wait for peers that + # never arrive. Reducing first fixes both, and makes the rebuild + # decision unanimous without a second collective. capacity = torch.tensor([tokens.shape[0]], dtype=torch.int64, device=tokens.device) dist.all_reduce(capacity, op=dist.ReduceOp.MAX, group=self.ep_group) - num_max_tokens_per_rank = int(capacity.item()) + # Rounded up so a sequence length that creeps upward does not rebuild + # on almost every step. + granularity = 512 + num_max_tokens_per_rank = -(-int(capacity.item()) // granularity) * granularity if self._deepep_exchange is None or num_max_tokens_per_rank > self._deepep_exchange.num_max_tokens_per_rank: if self._deepep_exchange is not None: - self._deepep_exchange.destroy() + # Kept rather than destroyed. A backward from an earlier + # forward replays against the buffer that forward dispatched + # on, so releasing it here breaks + # out1 = model(b1); out2 = model(b2); (out1 + out2).backward() + # where b2 outgrew b1's buffer. + self._retired_exchanges.append(self._deepep_exchange) self._deepep_exchange = DeepEPExchange( ep_group=self.ep_group, num_experts=self.num_experts, top_k=self.top_k, hidden_size=self.hidden_size, num_max_tokens_per_rank=num_max_tokens_per_rank, + num_sms=self.comm_num_sm, + qp_margin=self.comm_qp_margin, ) received, recv_weights, exchange = deepep_dispatch(self._deepep_exchange, tokens, ro.selected_experts, ro.top_scores) handle = exchange.last_handle - # Two quantities have to line up for the grouped GEMM, and they come - # from different places: the rows are the valid prefix of a buffer - # sized for the worst case, and the per-expert counts arrive as a - # prefix sum. Differencing the same prefix sum that bounds the rows - # keeps the two consistent. + # The routing weight is applied here in both score modes and is never + # handed to combine: DeepEP's combine does not multiply rows by the + # weights it carries, it transports and reduces them separately, so + # passing them there would drop them from the result entirely. + # + # Which side of the experts it goes on has to match the collective + # path, because the expert MLP is not linear and the two placements do + # not commute through SwiGLU. + weights = None if recv_weights is None else recv_weights.reshape(-1, 1) + if weights is not None and self.score_apply == "pre": + received = (received.float() * weights[:received.shape[0]]).to(received.dtype) + weights = None + + # The grouped GEMM needs per-expert row counts, which arrive as a + # prefix sum over the experts this rank owns. Differencing it recovers + # the counts. The rows themselves are not trimmed to the prefix: all + # three expert paths bound their work by these counts and ignore what + # lies beyond, so trimming would only add a device-to-host + # synchronisation in front of every layer's GEMM. 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())] if counts.numel() != self.num_local_experts: raise RuntimeError(f"DeepEP returned {counts.numel()} expert counts, but this rank owns " f"{self.num_local_experts} experts") - # The weights arrive in a buffer sized for the worst case, like the - # rows did, so they are trimmed to the same prefix. Combine reduces - # rows and weights together and rejects a pair whose lengths disagree. - if recv_weights is not None: - recv_weights = recv_weights[:received.shape[0]] - - # The router weights must be applied exactly once. In "pre" mode the - # NCCL path scales tokens before the experts see them, so the arrived - # rows are scaled the same way; in "post" mode the scaling is part of - # the reduction and is handed to combine instead. In expand mode each - # row is one token-expert slot, so a weight per row broadcasts across - # the hidden dimension. - if self.score_apply == "pre" and recv_weights is not None: - scale = recv_weights.reshape(-1, 1) - received = (received.to(torch.float32) * scale).to(received.dtype) - combine_weights = None - else: - combine_weights = recv_weights - expert_output = self.experts(received, counts) - return deepep_combine(exchange, expert_output, handle, combine_weights) + + if weights is not None: + expert_output = (expert_output.float() * weights[:expert_output.shape[0]]).to(expert_output.dtype) + + return deepep_combine(exchange, expert_output, handle) def forward( self, diff --git a/deepspeed/module_inject/auto_ep_presets/base.py b/deepspeed/module_inject/auto_ep_presets/base.py index 7f214210aed9..9f9814d4e35d 100644 --- a/deepspeed/module_inject/auto_ep_presets/base.py +++ b/deepspeed/module_inject/auto_ep_presets/base.py @@ -110,6 +110,9 @@ class AutoEPConfig: route_scale: float = 1.0 score_apply: Literal["auto", "pre", "post"] = "auto" combine_impl: Literal["auto", "weighted_sum", "legacy_bmm"] = "auto" + comm_backend: Literal["comm", "deepep"] = "comm" + comm_num_sm: int = 12 + comm_qp_margin: int = 4 num_expert_groups: int | None = None num_limited_groups: int | None = None score_func: Literal["auto", "softmax", "sigmoid"] = "auto" diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py index 0fda54b69cbc..aa4e94cc138d 100755 --- a/deepspeed/runtime/engine.py +++ b/deepspeed/runtime/engine.py @@ -824,6 +824,11 @@ def __del__(self): logger.debug("DeepSpeedEngine.__del__ cleanup skipped: %s", exc, exc_info=True) def destroy(self): + # DeepEP buffers ask the library not to reclaim them, so they outlive + # the engine unless something releases them here. + from deepspeed.module_inject.auto_ep_comm import destroy_all_exchanges + destroy_all_exchanges() + self._release_deepcompile_compiled_backward_state() self._release_deepcompile_dynamo_config() optimizer = getattr(self, "optimizer", None) diff --git a/docs/code-docs/source/autoep.rst b/docs/code-docs/source/autoep.rst index 29f05edd978d..8175f21146e7 100644 --- a/docs/code-docs/source/autoep.rst +++ b/docs/code-docs/source/autoep.rst @@ -87,24 +87,40 @@ Weights-only/module-only Universal Checkpoint loads use the converted **Communication backend (optional):** The expert AllToAll can be carried by `DeepEP `__ -instead of NCCL. This is opt-in and off by default; jobs that set nothing keep -the NCCL path unchanged. +instead of the default collectives. This is opt-in and off by default; jobs +that set nothing keep the existing path unchanged. -.. code-block:: bash +.. code-block:: json - export DEEPSPEED_AUTOEP_COMM_BACKEND=deepep # default: nccl - export DEEPSPEED_AUTOEP_COMM_SMS=12 # communication SM budget + { + "expert_parallel": { + "enabled": true, + "autoep_size": 8, + "comm_backend": "deepep", + "comm_num_sm": 12, + "comm_qp_margin": 4 + } + } + +- ``comm_backend``: ``"comm"`` (default) uses ``deepspeed.comm`` collectives; + ``"deepep"`` uses DeepEP's dispatch and combine kernels. +- ``comm_num_sm``: SMs given to communication. Default 12. +- ``comm_qp_margin``: RDMA queue pairs reserved beyond one per SM. Default 4. On 16 H100s across two nodes, replaying routing captured from real training, DeepEP reduced payload AllToAll time from roughly 100 ms to 48 ms per step, and whole training steps from 454 ms to 371 ms. The advantage grows with routing -imbalance: at the most skewed step measured, the NCCL path degraded to 116 ms -while DeepEP stayed flat. +imbalance: at the most skewed step measured, the collective path degraded to +116 ms while DeepEP stayed flat. + +``comm_num_sm`` matters because communication competes with the expert GEMM for +SMs. The default of 12 was chosen by measuring whole steps: a smaller budget +slows the collective itself, and a larger one slows everything else. +``comm_qp_margin`` exists because DeepEP's automatic queue-pair count assumes +it is alone on the fabric, which exhausts the queue pairs ZeRO and the +data-parallel groups have already claimed in a training step. -``DEEPSPEED_AUTOEP_COMM_SMS`` matters because communication competes with the -expert GEMM for SMs. The default of 12 was chosen by measuring whole steps: a -smaller budget slows the collective itself, and a larger one slows everything -else. Requirements and limits: +Requirements and limits: - The ``deep_ep`` package must be installed. It is imported only when this backend is selected, so installations without it are unaffected. @@ -112,8 +128,10 @@ else. Requirements and limits: version the transport is unavailable regardless of the network. - DeepEP v1 (the legacy ``Buffer`` API, using NVSHMEM and IBGDA) is not supported. -- The backend is disabled, with a warning, on layers using folded tensor - parallelism (``tp_size > 1``); those layers fall back to NCCL. +- bfloat16 only. DeepEP's dispatch kernel does not handle fp16, and selecting + this backend for an fp16 run is rejected rather than silently downgraded. +- Not compatible with folded tensor parallelism + (``expert_tensor_parallel_size > 1``), which is rejected at setup. **Constraints:** diff --git a/tests/unit/module_inject/test_auto_ep_comm.py b/tests/unit/module_inject/test_auto_ep_comm.py index 9cd80e6854bb..055708801532 100644 --- a/tests/unit/module_inject/test_auto_ep_comm.py +++ b/tests/unit/module_inject/test_auto_ep_comm.py @@ -3,7 +3,6 @@ import ast import inspect -import os import textwrap import sys import unittest @@ -11,45 +10,78 @@ import torch from unittest import mock -from deepspeed.module_inject.auto_ep_comm import (AVAILABLE_BACKENDS, DEEPEP_BACKEND, NCCL_BACKEND, _conform_rows, - _DeepEPCombine, _DeepEPDispatch, _import_deep_ep, configured_backend, - configured_num_sms) +from deepspeed.module_inject.auto_ep_comm import (COMM_BACKEND, DEEPEP_BACKEND, SUPPORTED_DTYPES, _conform_rows, + _DeepEPCombine, _DeepEPDispatch, _import_deep_ep, _qps_for_sms, + assert_dtype_supported) +from deepspeed.module_inject.auto_ep_config import parse_autoep_config, validate_autoep_config class TestAutoEPCommBackendSelection(unittest.TestCase): + """Backend choice lives in the config, alongside the rest of AutoEP.""" - def test_defaults_to_nccl(self): - # An unset variable must leave existing jobs on the shipped path. - with mock.patch.dict(os.environ, {}, clear=True): - self.assertEqual(configured_backend(), NCCL_BACKEND) - self.assertEqual(configured_num_sms(), 0) + @staticmethod + def validate(config): + validate_autoep_config(config, world_size=1, pp_size=1, tp_size=1, sp_size=1) + + def test_defaults_to_the_collective_path(self): + # A config that asks for nothing must leave existing jobs unchanged. + config = parse_autoep_config({"enabled": True}) + + self.assertEqual(config.comm_backend, COMM_BACKEND) + self.assertEqual(config.comm_num_sm, 12) + self.assertEqual(config.comm_qp_margin, 4) def test_selects_deepep(self): - with mock.patch.dict(os.environ, {"DEEPSPEED_AUTOEP_COMM_BACKEND": "DeepEP"}): - self.assertEqual(configured_backend(), DEEPEP_BACKEND) + config = parse_autoep_config({"enabled": True, "comm_backend": "deepep"}) + + self.assertEqual(config.comm_backend, DEEPEP_BACKEND) + + def test_sm_budget_and_qp_margin_are_configurable(self): + config = parse_autoep_config({"enabled": True, "comm_num_sm": 24, "comm_qp_margin": 8}) + + self.assertEqual(config.comm_num_sm, 24) + self.assertEqual(config.comm_qp_margin, 8) def test_rejects_unknown_backend(self): # Failing loudly beats silently running the wrong transport. - with mock.patch.dict(os.environ, {"DEEPSPEED_AUTOEP_COMM_BACKEND": "moonep"}): - with self.assertRaises(ValueError) as caught: - configured_backend() - self.assertIn(str(AVAILABLE_BACKENDS), str(caught.exception)) + config = parse_autoep_config({"enabled": True, "comm_backend": "moonep"}) + + with self.assertRaises(ValueError) as caught: + self.validate(config) + self.assertIn("moonep", str(caught.exception)) + + def test_rejects_a_zero_sm_budget(self): + # Zero would hand the whole GPU to the collective. + config = parse_autoep_config({"enabled": True, "comm_num_sm": 0}) - def test_sm_budget_is_configurable(self): - with mock.patch.dict(os.environ, {"DEEPSPEED_AUTOEP_COMM_SMS": "24"}): - self.assertEqual(configured_num_sms(), 24) + with self.assertRaises(ValueError): + self.validate(config) - def test_blank_value_is_treated_as_unset(self): - # An empty variable is what an unset shell variable expands to. - with mock.patch.dict(os.environ, {"DEEPSPEED_AUTOEP_COMM_BACKEND": " "}): - self.assertEqual(configured_backend(), NCCL_BACKEND) - with mock.patch.dict(os.environ, {"DEEPSPEED_AUTOEP_COMM_SMS": ""}): - self.assertEqual(configured_num_sms(), 0) + def test_rejects_a_negative_qp_margin(self): + config = parse_autoep_config({"enabled": True, "comm_qp_margin": -1}) - def test_non_numeric_sm_budget_is_rejected(self): - with mock.patch.dict(os.environ, {"DEEPSPEED_AUTOEP_COMM_SMS": "many"}): - with self.assertRaises(ValueError): - configured_num_sms() + with self.assertRaises(ValueError): + self.validate(config) + + def test_queue_pairs_leave_room_for_the_control_path(self): + # DeepEP's own default exhausts the QPs that ZeRO and the + # data-parallel groups have already claimed in a training step. + self.assertEqual(_qps_for_sms(12, 4), 16) + + +class TestDtypeGuard(unittest.TestCase): + """DeepEP's dispatch kernel handles bfloat16, not fp16.""" + + def test_rejects_fp16(self): + with self.assertRaises(TypeError) as caught: + assert_dtype_supported(torch.float16) + self.assertIn("bfloat16", str(caught.exception)) + + def test_accepts_bfloat16(self): + self.assertIsNone(assert_dtype_supported(torch.bfloat16)) + + def test_fp16_is_not_quietly_in_the_supported_set(self): + self.assertNotIn(torch.float16, SUPPORTED_DTYPES) class TestDeepEPPreflight(unittest.TestCase): @@ -61,7 +93,7 @@ def test_missing_package_names_its_requirements(self): _import_deep_ep() message = str(caught.exception) self.assertIn("2.30.4", message) - self.assertIn("DEEPSPEED_AUTOEP_COMM_BACKEND", message) + self.assertIn("comm_backend", message) def test_old_torch_nccl_warns_but_does_not_block(self): # torch reports the NCCL it bundles, which DeepEP need not be using; @@ -126,45 +158,54 @@ def test_conforms_a_one_dimensional_weight_buffer(self): self.assertEqual(tuple(_conform_rows(grad, (12, )).shape), (12, )) -class TestRoutedRowAgreement(unittest.TestCase): - """Rows and their weights must describe the same arrivals. +class TestRoutingWeightsAreApplied(unittest.TestCase): + """The layer must apply the routing weights itself. - DeepEP hands back whole worst-case buffers, and the layer trims the rows to - the arrivals the prefix sum reports. Leaving the weights at full length - makes combine reduce a row count and a weight count that disagree, which is - silent in "pre" mode and fatal in "post" mode. + DeepEP's combine does not multiply rows by the topk_weights it is given. + It transports and reduces them separately and returns them as a second + output, so handing the weights to combine drops them from the result: the + expert outputs come back summed but unweighted, which trains quietly and + wrongly. """ BUFFER_ROWS = 32 - ARRIVED_ROWS = 12 HIDDEN = 8 LOCAL_EXPERTS = 4 + WEIGHT = 0.25 def route(self, score_apply): - """Drive _deepep_route with a stub exchange, returning combine's inputs.""" + """Drive _deepep_route with a stub exchange, recording what each stage saw.""" from deepspeed.module_inject import auto_ep_layer - prefix = torch.tensor([3, 6, 9, self.ARRIVED_ROWS], dtype=torch.int64) + prefix = torch.tensor([3, 6, 9, 12], dtype=torch.int64) handle = mock.Mock(psum_num_recv_tokens_per_expert=prefix) - exchange = mock.Mock(last_handle=handle, num_max_tokens_per_rank=self.BUFFER_ROWS) + exchange = mock.Mock(last_handle=handle, num_max_tokens_per_rank=1024) buffer_rows = torch.ones((self.BUFFER_ROWS, self.HIDDEN)) - buffer_weights = torch.ones(self.BUFFER_ROWS) + buffer_weights = torch.full((self.BUFFER_ROWS, ), self.WEIGHT) seen = {} def fake_dispatch(_exchange, *_args): return buffer_rows, buffer_weights, exchange - def fake_combine(_exchange, rows, _handle, topk_weights=None): - seen["rows"] = rows - seen["weights"] = topk_weights + def fake_combine(_exchange, rows, _handle, **kwargs): + seen["combine_rows"] = rows + seen["combine_kwargs"] = kwargs + return rows + + def fake_experts(rows, counts): + seen["expert_input"] = rows + seen["counts"] = counts return rows layer = mock.Mock( _deepep_exchange=exchange, + _retired_exchanges=[], num_local_experts=self.LOCAL_EXPERTS, score_apply=score_apply, - experts=lambda rows, counts: rows, + comm_num_sm=12, + comm_qp_margin=4, + experts=fake_experts, ) with mock.patch.object(auto_ep_layer, "deepep_dispatch", fake_dispatch), \ @@ -175,20 +216,67 @@ def fake_combine(_exchange, rows, _handle, topk_weights=None): selected_experts=torch.zeros((4, 2), dtype=torch.long), num_tokens_per_expert=torch.zeros(self.LOCAL_EXPERTS, dtype=torch.long), ) - auto_ep_layer.AutoEPMoELayer._deepep_route(layer, torch.ones((4, self.HIDDEN)), router_output) + tokens = torch.ones((4, self.HIDDEN), dtype=torch.bfloat16) + seen["result"] = auto_ep_layer.AutoEPMoELayer._deepep_route(layer, tokens, router_output) return seen - def test_post_mode_hands_combine_one_weight_per_row(self): + def test_combine_is_never_given_the_weights(self): + for score_apply in ("pre", "post"): + with self.subTest(score_apply=score_apply): + seen = self.route(score_apply) + + self.assertNotIn("topk_weights", seen["combine_kwargs"]) + + def test_post_mode_weights_the_expert_output(self): seen = self.route("post") - self.assertEqual(seen["rows"].shape[0], self.ARRIVED_ROWS) - self.assertEqual(seen["weights"].shape[0], self.ARRIVED_ROWS) + # The experts saw unscaled rows, and the scaling landed after them. + self.assertTrue(torch.allclose(seen["expert_input"].float(), torch.ones(1))) + self.assertTrue(torch.allclose(seen["combine_rows"].float(), torch.full((1, ), self.WEIGHT))) - def test_pre_mode_scales_only_the_rows_that_arrived(self): + def test_pre_mode_weights_the_expert_input(self): seen = self.route("pre") - self.assertEqual(seen["rows"].shape[0], self.ARRIVED_ROWS) - self.assertIsNone(seen["weights"], "pre mode applies the weights before the experts, not during combine") + # Scaling landed before the experts, and is not applied a second time. + self.assertTrue(torch.allclose(seen["expert_input"].float(), torch.full((1, ), self.WEIGHT))) + self.assertTrue(torch.allclose(seen["combine_rows"].float(), torch.full((1, ), self.WEIGHT))) + + def test_rows_are_not_trimmed_before_the_grouped_gemm(self): + # Trimming needs prefix[-1] on the host, which drains the pipeline in + # front of every layer's GEMM. The expert paths bound their own work by + # the counts, so the untrimmed buffer is safe to pass straight through. + seen = self.route("post") + + self.assertEqual(seen["expert_input"].shape[0], self.BUFFER_ROWS) + self.assertTrue(torch.equal(seen["counts"], torch.tensor([3, 3, 3, 3], dtype=torch.int32))) + + +class TestBufferLifecycle(unittest.TestCase): + """A grown buffer must not invalidate a backward that is still pending.""" + + def test_outgrown_buffer_is_retired_rather_than_destroyed(self): + from deepspeed.module_inject import auto_ep_layer + + old = mock.Mock(num_max_tokens_per_rank=512) + layer = mock.Mock(_deepep_exchange=old, _retired_exchanges=[], comm_num_sm=12, comm_qp_margin=4) + + with mock.patch.object(auto_ep_layer, "DeepEPExchange") as built, \ + mock.patch.object(auto_ep_layer, "deepep_dispatch", side_effect=RuntimeError("stop here")), \ + mock.patch.object(auto_ep_layer.dist, "all_reduce", lambda *a, **k: None): + router_output = auto_ep_layer.RouterOutput( + top_scores=torch.ones((600, 1)), + selected_experts=torch.zeros((600, 1), dtype=torch.long), + num_tokens_per_expert=torch.zeros(4, dtype=torch.long), + ) + with self.assertRaises(RuntimeError): + auto_ep_layer.AutoEPMoELayer._deepep_route(layer, torch.ones((600, 8), dtype=torch.bfloat16), + router_output) + + old.destroy.assert_not_called() + self.assertIn(old, layer._retired_exchanges) + # Capacity is rounded up so a slowly growing sequence length does not + # rebuild on nearly every step. + self.assertEqual(built.call_args.kwargs["num_max_tokens_per_rank"], 1024) class TestAutogradSignatures(unittest.TestCase): @@ -220,9 +308,11 @@ def test_dispatch_backward_returns_a_gradient_per_input(self): self.assertEqual(self.gradient_count(_DeepEPDispatch.backward), inputs) def test_combine_backward_returns_a_gradient_per_input(self): + # Combine takes no weights: DeepEP would not apply them, so the layer + # folds them into the rows before calling it. inputs = len(inspect.signature(_DeepEPCombine.forward).parameters) - 1 - self.assertEqual(inputs, 4) + self.assertEqual(inputs, 3) self.assertEqual(self.gradient_count(_DeepEPCombine.backward), inputs) def test_dispatch_forward_returns_weights_so_they_stay_differentiable(self): From 615461f3a34b9c9090e4c63e66fe9d30b30afa50 Mon Sep 17 00:00:00 2001 From: yh0903 Date: Wed, 5 Aug 2026 11:50:01 -0700 Subject: [PATCH 6/7] Cut dispatch buffers to the arrivals the handle records 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 --- deepspeed/module_inject/auto_ep_layer.py | 27 ++++++++++------- tests/unit/module_inject/test_auto_ep_comm.py | 29 ++++++++++++------- 2 files changed, 34 insertions(+), 22 deletions(-) diff --git a/deepspeed/module_inject/auto_ep_layer.py b/deepspeed/module_inject/auto_ep_layer.py index bf0c64226597..43808152a07a 100644 --- a/deepspeed/module_inject/auto_ep_layer.py +++ b/deepspeed/module_inject/auto_ep_layer.py @@ -569,10 +569,9 @@ def set_deepspeed_parallelism( # returns token-major rows cannot satisfy. Refusing beats # falling back, which leaves the job reporting a backend it is # not running and the speedup unexplained. - raise ValueError( - f'comm_backend="{DEEPEP_BACKEND}" does not support folded tensor parallelism ' - f"(expert_tensor_parallel_size={folding_group_handles.spec.tp_size}). Set " - 'expert_tensor_parallel_size to 1, or comm_backend to "comm".') + raise ValueError(f'comm_backend="{DEEPEP_BACKEND}" does not support folded tensor parallelism ' + f"(expert_tensor_parallel_size={folding_group_handles.spec.tp_size}). Set " + 'expert_tensor_parallel_size to 1, or comm_backend to "comm".') self.ep_group_name = folding_group_handles.ep_group_name self.ep_group = folding_group_handles.ep_group self.tp_group = folding_group_handles.tp_group @@ -648,6 +647,15 @@ def _deepep_route(self, tokens: torch.Tensor, ro: "RouterOutput") -> torch.Tenso ro.top_scores) handle = exchange.last_handle + # Dispatch returns buffers sized for the worst case, and combine reads + # exactly the rows the handle says arrived, so everything downstream is + # cut to that length before it is used. The count is taken from the + # handle, where it is already a Python int, rather than off the end of + # the prefix sum on the device: reading it there would put a + # device-to-host synchronisation in front of every layer's expert GEMM. + arrived = handle.num_expanded_tokens + received = received[:arrived] + # The routing weight is applied here in both score modes and is never # handed to combine: DeepEP's combine does not multiply rows by the # weights it carries, it transports and reduces them separately, so @@ -656,17 +664,14 @@ def _deepep_route(self, tokens: torch.Tensor, ro: "RouterOutput") -> torch.Tenso # Which side of the experts it goes on has to match the collective # path, because the expert MLP is not linear and the two placements do # not commute through SwiGLU. - weights = None if recv_weights is None else recv_weights.reshape(-1, 1) + weights = None if recv_weights is None else recv_weights[:arrived].reshape(-1, 1) if weights is not None and self.score_apply == "pre": - received = (received.float() * weights[:received.shape[0]]).to(received.dtype) + received = (received.float() * weights).to(received.dtype) weights = None # The grouped GEMM needs per-expert row counts, which arrive as a # prefix sum over the experts this rank owns. Differencing it recovers - # the counts. The rows themselves are not trimmed to the prefix: all - # three expert paths bound their work by these counts and ignore what - # lies beyond, so trimming would only add a device-to-host - # synchronisation in front of every layer's GEMM. + # the counts. prefix = handle.psum_num_recv_tokens_per_expert counts = torch.diff(prefix, prepend=prefix.new_zeros(1)).to(torch.int32) if counts.numel() != self.num_local_experts: @@ -676,7 +681,7 @@ def _deepep_route(self, tokens: torch.Tensor, ro: "RouterOutput") -> torch.Tenso expert_output = self.experts(received, counts) if weights is not None: - expert_output = (expert_output.float() * weights[:expert_output.shape[0]]).to(expert_output.dtype) + expert_output = (expert_output.float() * weights).to(expert_output.dtype) return deepep_combine(exchange, expert_output, handle) diff --git a/tests/unit/module_inject/test_auto_ep_comm.py b/tests/unit/module_inject/test_auto_ep_comm.py index 055708801532..f72aac9bdeb0 100644 --- a/tests/unit/module_inject/test_auto_ep_comm.py +++ b/tests/unit/module_inject/test_auto_ep_comm.py @@ -13,6 +13,7 @@ from deepspeed.module_inject.auto_ep_comm import (COMM_BACKEND, DEEPEP_BACKEND, SUPPORTED_DTYPES, _conform_rows, _DeepEPCombine, _DeepEPDispatch, _import_deep_ep, _qps_for_sms, assert_dtype_supported) +from deepspeed.module_inject import auto_ep_layer from deepspeed.module_inject.auto_ep_config import parse_autoep_config, validate_autoep_config @@ -169,16 +170,15 @@ class TestRoutingWeightsAreApplied(unittest.TestCase): """ BUFFER_ROWS = 32 + ARRIVED_ROWS = 12 HIDDEN = 8 LOCAL_EXPERTS = 4 WEIGHT = 0.25 def route(self, score_apply): """Drive _deepep_route with a stub exchange, recording what each stage saw.""" - from deepspeed.module_inject import auto_ep_layer - - prefix = torch.tensor([3, 6, 9, 12], dtype=torch.int64) - handle = mock.Mock(psum_num_recv_tokens_per_expert=prefix) + prefix = torch.tensor([3, 6, 9, self.ARRIVED_ROWS], dtype=torch.int64) + handle = mock.Mock(psum_num_recv_tokens_per_expert=prefix, num_expanded_tokens=self.ARRIVED_ROWS) exchange = mock.Mock(last_handle=handle, num_max_tokens_per_rank=1024) buffer_rows = torch.ones((self.BUFFER_ROWS, self.HIDDEN)) @@ -241,22 +241,29 @@ def test_pre_mode_weights_the_expert_input(self): self.assertTrue(torch.allclose(seen["expert_input"].float(), torch.full((1, ), self.WEIGHT))) self.assertTrue(torch.allclose(seen["combine_rows"].float(), torch.full((1, ), self.WEIGHT))) - def test_rows_are_not_trimmed_before_the_grouped_gemm(self): - # Trimming needs prefix[-1] on the host, which drains the pipeline in - # front of every layer's GEMM. The expert paths bound their own work by - # the counts, so the untrimmed buffer is safe to pass straight through. + def test_combine_receives_exactly_the_rows_that_arrived(self): + # Dispatch returns a worst-case buffer while combine reads the rows the + # handle recorded. Handing it the whole buffer reads past what the + # handle describes, which does not raise: it faults. seen = self.route("post") - self.assertEqual(seen["expert_input"].shape[0], self.BUFFER_ROWS) + self.assertEqual(seen["expert_input"].shape[0], self.ARRIVED_ROWS) + self.assertEqual(seen["combine_rows"].shape[0], self.ARRIVED_ROWS) self.assertTrue(torch.equal(seen["counts"], torch.tensor([3, 3, 3, 3], dtype=torch.int32))) + def test_row_count_comes_from_the_handle_not_the_device(self): + # Reading it off the prefix sum needs a device-to-host synchronisation + # in front of every layer's GEMM; the handle already holds it as an int. + source = inspect.getsource(auto_ep_layer.AutoEPMoELayer._deepep_route) + + self.assertIn("arrived = handle.num_expanded_tokens", source) + self.assertNotIn("psum_num_recv_tokens_per_expert[-1]", source) + class TestBufferLifecycle(unittest.TestCase): """A grown buffer must not invalidate a backward that is still pending.""" def test_outgrown_buffer_is_retired_rather_than_destroyed(self): - from deepspeed.module_inject import auto_ep_layer - old = mock.Mock(num_max_tokens_per_rank=512) layer = mock.Mock(_deepep_exchange=old, _retired_exchanges=[], comm_num_sm=12, comm_qp_margin=4) From ab4e7ccbf78869f869208172227d7e20f1f6902c Mon Sep 17 00:00:00 2001 From: yh0903 Date: Wed, 5 Aug 2026 13:48:02 -0700 Subject: [PATCH 7/7] Agree the DeepEP buffer size once per shape rather than per layer 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 --- deepspeed/module_inject/auto_ep_comm.py | 6 +++ deepspeed/module_inject/auto_ep_layer.py | 38 +++++++++++-------- docs/code-docs/source/autoep.rst | 9 +++-- tests/unit/module_inject/test_auto_ep_comm.py | 23 ++++++++++- 4 files changed, 55 insertions(+), 21 deletions(-) diff --git a/deepspeed/module_inject/auto_ep_comm.py b/deepspeed/module_inject/auto_ep_comm.py index 86645c4d5928..8100b0f7bdd7 100644 --- a/deepspeed/module_inject/auto_ep_comm.py +++ b/deepspeed/module_inject/auto_ep_comm.py @@ -182,6 +182,12 @@ def __init__(self, # on its own. Registering it means a process that never calls the # layer's teardown can still release every buffer in one call. _LIVE_EXCHANGES.append(self) + # Buffer construction is collective and allocates fabric resources, so + # it is where an unsuitable cluster fails, often by killing the process + # without raising. Recording each one lets a post-mortem tell a buffer + # that was never built from one that was built and then used. + logger.info(f"AutoEP DeepEP buffer {len(_LIVE_EXCHANGES)} built: " + f"capacity={num_max_tokens_per_rank} sms={num_sms} qps={_qps_for_sms(num_sms, qp_margin)}") def dispatch(self, tokens: torch.Tensor, topk_idx: torch.Tensor, topk_weights: torch.Tensor): """Send tokens to their experts, returning rows, weights and handle. diff --git a/deepspeed/module_inject/auto_ep_layer.py b/deepspeed/module_inject/auto_ep_layer.py index 43808152a07a..67032c72d8cb 100644 --- a/deepspeed/module_inject/auto_ep_layer.py +++ b/deepspeed/module_inject/auto_ep_layer.py @@ -528,6 +528,10 @@ def __init__( self.comm_num_sm = config.comm_num_sm self.comm_qp_margin = config.comm_qp_margin self._deepep_exchange = None + # The agreed capacity and the shape it was agreed for, so the + # agreement is not repeated on every layer of every step. + self._deepep_capacity = 0 + self._deepep_tokens_per_rank = -1 # Buffers outgrown by a larger batch. They stay alive rather than being # released, because a backward still pending from an earlier forward # replays against the buffer that forward used. Capacity is rounded up @@ -611,21 +615,23 @@ def _deepep_route(self, tokens: torch.Tensor, ro: "RouterOutput") -> torch.Tenso # batch is larger: variable sequence lengths or a small warm-up batch # would otherwise exceed a capacity fixed by the first call. # - # Capacity is agreed across the group before it is used. Ranks hold - # different local token counts, and DeepEP requires the same - # num_max_tokens_per_rank everywhere; sizing from the local count also - # has some ranks entering construction while others do not, and - # construction is collective, so those that entered wait for peers that - # never arrive. Reducing first fixes both, and makes the rebuild - # decision unanimous without a second collective. - capacity = torch.tensor([tokens.shape[0]], dtype=torch.int64, device=tokens.device) - dist.all_reduce(capacity, op=dist.ReduceOp.MAX, group=self.ep_group) - # Rounded up so a sequence length that creeps upward does not rebuild - # on almost every step. - granularity = 512 - num_max_tokens_per_rank = -(-int(capacity.item()) // granularity) * granularity - - if self._deepep_exchange is None or num_max_tokens_per_rank > self._deepep_exchange.num_max_tokens_per_rank: + # DeepEP requires the same num_max_tokens_per_rank on every rank, and + # construction is collective, so the size is agreed rather than taken + # from the local count. The agreement is cached against the shape that + # produced it: ranks in an expert-parallel group are handed the same + # 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: + capacity = torch.tensor([tokens.shape[0]], dtype=torch.int64, device=tokens.device) + dist.all_reduce(capacity, op=dist.ReduceOp.MAX, group=self.ep_group) + # Rounded up so a sequence length that creeps upward does not + # rebuild on almost every step. + granularity = 512 + self._deepep_capacity = -(-int(capacity.item()) // granularity) * granularity + self._deepep_tokens_per_rank = tokens.shape[0] + + if self._deepep_exchange is None or self._deepep_capacity > self._deepep_exchange.num_max_tokens_per_rank: if self._deepep_exchange is not None: # Kept rather than destroyed. A backward from an earlier # forward replays against the buffer that forward dispatched @@ -638,7 +644,7 @@ def _deepep_route(self, tokens: torch.Tensor, ro: "RouterOutput") -> torch.Tenso num_experts=self.num_experts, top_k=self.top_k, hidden_size=self.hidden_size, - num_max_tokens_per_rank=num_max_tokens_per_rank, + num_max_tokens_per_rank=self._deepep_capacity, num_sms=self.comm_num_sm, qp_margin=self.comm_qp_margin, ) diff --git a/docs/code-docs/source/autoep.rst b/docs/code-docs/source/autoep.rst index 8175f21146e7..2a7d909a7dd9 100644 --- a/docs/code-docs/source/autoep.rst +++ b/docs/code-docs/source/autoep.rst @@ -108,10 +108,11 @@ that set nothing keep the existing path unchanged. - ``comm_qp_margin``: RDMA queue pairs reserved beyond one per SM. Default 4. On 16 H100s across two nodes, replaying routing captured from real training, -DeepEP reduced payload AllToAll time from roughly 100 ms to 48 ms per step, and -whole training steps from 454 ms to 371 ms. The advantage grows with routing -imbalance: at the most skewed step measured, the collective path degraded to -116 ms while DeepEP stayed flat. +DeepEP reduced payload AllToAll time from roughly 100 ms to 48 ms per step. A +full SFT step on Qwen3.5-MoE went from 360.6 ms to 310.6 ms, a 1.16x speedup +that removes 13.9% of the step. The advantage grows with routing imbalance: at +the most skewed step measured, the collective path degraded to 116 ms while +DeepEP stayed flat. ``comm_num_sm`` matters because communication competes with the expert GEMM for SMs. The default of 12 was chosen by measuring whole steps: a smaller budget diff --git a/tests/unit/module_inject/test_auto_ep_comm.py b/tests/unit/module_inject/test_auto_ep_comm.py index f72aac9bdeb0..3d314e93c37f 100644 --- a/tests/unit/module_inject/test_auto_ep_comm.py +++ b/tests/unit/module_inject/test_auto_ep_comm.py @@ -200,6 +200,8 @@ def fake_experts(rows, counts): layer = mock.Mock( _deepep_exchange=exchange, + _deepep_capacity=1024, + _deepep_tokens_per_rank=4, _retired_exchanges=[], num_local_experts=self.LOCAL_EXPERTS, score_apply=score_apply, @@ -263,9 +265,28 @@ def test_row_count_comes_from_the_handle_not_the_device(self): class TestBufferLifecycle(unittest.TestCase): """A grown buffer must not invalidate a backward that is still pending.""" + def test_capacity_agreement_is_skipped_when_the_shape_is_unchanged(self): + # The agreement is collective and ends in a device-to-host read, so + # repeating it per layer per step puts a synchronisation in the middle + # of the forward pass. + source = inspect.getsource(auto_ep_layer.AutoEPMoELayer._deepep_route) + tree = ast.parse(textwrap.dedent(source)) + guards = [ + node for node in ast.walk(tree) + if isinstance(node, ast.If) and "_deepep_tokens_per_rank" in ast.dump(node.test) + ] + + self.assertTrue(guards, "the all_reduce must be guarded by an unchanged-shape check") + self.assertIn("all_reduce", ast.dump(guards[0])) + def test_outgrown_buffer_is_retired_rather_than_destroyed(self): old = mock.Mock(num_max_tokens_per_rank=512) - layer = mock.Mock(_deepep_exchange=old, _retired_exchanges=[], comm_num_sm=12, comm_qp_margin=4) + layer = mock.Mock(_deepep_exchange=old, + _deepep_capacity=512, + _deepep_tokens_per_rank=-1, + _retired_exchanges=[], + comm_num_sm=12, + comm_qp_margin=4) with mock.patch.object(auto_ep_layer, "DeepEPExchange") as built, \ mock.patch.object(auto_ep_layer, "deepep_dispatch", side_effect=RuntimeError("stop here")), \