Add an opt-in DeepEP transport for the AutoEP expert all-to-all

deepspeed/module_inject/auto_ep_config.py

@@ -58,6 +58,10 @@ 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.comm_max_tokens_per_rank = param_dict.get("comm_max_tokens_per_rank", 0)
     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 +161,24 @@ 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}")
+    # Zero means "size from the first batch"; a negative size is meaningless.
+    if not isinstance(config.comm_max_tokens_per_rank, int) or config.comm_max_tokens_per_rank < 0:
+        raise ValueError("comm_max_tokens_per_rank must be a non-negative integer, "
+                         f"got {config.comm_max_tokens_per_rank!r}")
+
     # Validate score_func
     valid_score_func = ("auto", "softmax", "sigmoid")
     if config.score_func not in valid_score_func:

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, 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
 from deepspeed.moe.ep_experts import GroupedExperts
@@ -520,6 +522,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.
+        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
+        # The agreed capacity and the shape it was agreed for, so the
+        # agreement is not repeated on every layer of every step.
+        self.comm_max_tokens_per_rank = config.comm_max_tokens_per_rank
         self._register_logit_hook()
 
     def _register_logit_hook(self):
@@ -550,6 +561,15 @@ 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. 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
@@ -572,6 +592,100 @@ 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
+        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)
+
+        # Sized once, on the first forward, and then fixed.
+        #
+        # DeepEP requires the same num_max_tokens_per_rank on every rank, and
+        # construction is collective, so the size cannot come from the local
+        # token count. It is agreed with one all-reduce, which is safe here
+        # because every rank reaches its first forward together.
+        #
+        # It deliberately does not re-agree later. A rank cannot tell whether
+        # its peers also need more room without asking them, and asking only
+        # when the local count grew makes the collective conditional on local
+        # data: if the counts ever diverge, some ranks enter the all-reduce and
+        # the rest do not, and the job hangs. Growing on demand would also cost
+        # a device-to-host read in the middle of every layer of every step,
+        # measured at roughly a fifth of the step. So a batch that outgrows the
+        # buffer is a configuration error with a named remedy, not a stall.
+        if self._deepep_exchange is None:
+            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 varies a little still fits.
+            granularity = 512
+            sized_for = self.comm_max_tokens_per_rank or -(-int(capacity.item()) // granularity) * granularity
+            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=sized_for,
+                num_sms=self.comm_num_sm,
+                qp_margin=self.comm_qp_margin,
+            )
+
+        if tokens.shape[0] > self._deepep_exchange.num_max_tokens_per_rank:
+            raise RuntimeError(
+                f"this rank routed {tokens.shape[0]} tokens, more than the "
+                f"{self._deepep_exchange.num_max_tokens_per_rank} the DeepEP buffer was built for. The buffer is "
+                "sized once, from the first batch, because resizing is collective and cannot be decided from one "
+                "rank's token count. Set comm_max_tokens_per_rank in the expert_parallel config to the largest "
+                "per-rank token count this job will produce, which is micro_batch_size * seq_len when sequences "
+                'are padded to a fixed length, or set comm_backend="comm".')
+
+        received, recv_weights, exchange = deepep_dispatch(self._deepep_exchange, tokens, ro.selected_experts,
+                                                           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
+        # 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[:arrived].reshape(-1, 1)
+        if weights is not None and self.score_apply == "pre":
+            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.
+        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:
+            raise RuntimeError(f"DeepEP returned {counts.numel()} expert counts, but this rank owns "
+                               f"{self.num_local_experts} experts")
+
+        expert_output = self.experts(received, counts)
+
+        if weights is not None:
+            expert_output = (expert_output.float() * weights).to(expert_output.dtype)
+
+        return deepep_combine(exchange, expert_output, handle)
+
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -601,6 +715,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 (
@@ -678,21 +795,31 @@ def forward(
                     num_tokens_per_expert=ro.num_tokens_per_expert,
                 )
 
-            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)
+                deepep_combined = True
+            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,
                                       restore_ctx,
                                       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 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(
                 expert_output,

deepspeed/module_inject/auto_ep_presets/base.py

@@ -110,6 +110,10 @@ 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
+    comm_max_tokens_per_rank: int = 0
     num_expert_groups: int | None = None
     num_limited_groups: int | None = None
     score_func: Literal["auto", "softmax", "sigmoid"] = "auto"

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)

docs/code-docs/source/autoep.rst

@@ -84,6 +84,66 @@ 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 <https://github.com/deepseek-ai/DeepEP>`__
+instead of the default collectives. This is opt-in and off by default; jobs
+that set nothing keep the existing path unchanged.
+
+.. code-block:: json
+
+    {
+      "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.
+- ``comm_max_tokens_per_rank``: largest per-rank token count the job will
+  produce, which is ``micro_batch_size * seq_len`` when sequences are padded to
+  a fixed length. Default 0, which sizes the buffer from the first batch. Set
+  it when later batches can be larger than the first: the buffer is sized once
+  and a batch that outgrows it is an error, because resizing is collective and
+  cannot be decided from one rank's token count without risking a hang.
+
+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. A
+full SFT step on Qwen3.5-MoE went from roughly 325 ms to 266 ms, a 1.2x speedup
+that removes about 18% of the step, reproduced across two independent jobs
+(1.21x and 1.24x). Both backends are measured in the same job, on the same pods
+and alternating, since the same measurement varied by a quarter between jobs;
+the figures are medians rather than single observations, and DeepEP's own
+median moved by 0.3% between the two jobs while the collective baseline moved
+by 2.5%. 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: 8 SMs gave a median
+297.9 ms against 265.4 ms at 12, and larger budgets were slower again.
+``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.
+
+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.
+- 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:**
 
 - ``autoep_size`` must divide ``num_experts`` for all detected MoE layers.