AutoEP (Automatic Expert Parallelism)

AutoEP automatically detects MoE layers in Hugging Face models and replaces them with EP-enabled versions, requiring zero model code changes. It follows the pattern of AutoTP (Automatic Tensor Parallelism).

This API is separate from the explicit deepspeed.moe.layer.MoE layer API. For the explicit DeepSpeed MoE layer API, see Mixture of Experts (DeepSpeed MoE).

Built-in AutoEP presets: mixtral (Mixtral), qwen3_moe (Qwen3-MoE), qwen3_5_moe (Qwen3.5-MoE), deepseek_v2 (DeepSeek-V2), deepseek_v3 (DeepSeek-V3), and minimax_m3 (MiniMax-M3).

The preset name means AutoEP knows the router, expert, and weight naming patterns for that model family. Running a Hugging Face model also requires a Transformers build that exposes the matching config/model classes, model.config.model_type value, and fused expert layout.

AutoEP preset compatibility by Transformers version

Preset

Minimum Transformers version

Notes

mixtral

5.0.0

qwen3_moe

5.0.0

Also covers Qwen2-MoE when the installed Transformers build uses the validated fused expert layout. Qwen3-MoE classes appear in 4.51.3, but the tested 4.x builds do not match the validated AutoEP layout.

qwen3_5_moe

5.2.0

Requires the Qwen3.5 text-backbone qwen3_5_moe_text model type; for performance on Qwen3.5’s Gated DeltaNet layers, install optimized kernels. See the Hugging Face Transformers kernel loading docs and the Qwen FlashQLA blog.

deepseek_v2

5.0.0

load_balance_coeff / expert-bias auxiliary-loss-free load balancing is not currently supported; non-null values are rejected.

deepseek_v3

5.0.0

load_balance_coeff / expert-bias auxiliary-loss-free load balancing is not currently supported; non-null values are rejected.

minimax_m3

5.15.0

Requires the MiniMax-M3 text-backbone minimax_m3_vl_text model type. The expert MLP uses the clamped GPT-OSS activation (swiglu_oai), selected by the preset. load_balance_coeff / expert-bias auxiliary-loss-free load balancing is not currently supported; non-null values are rejected.

ZeRO compatibility: Stages 0, 1, and 2, plus constrained Stage 3 support. Stage 3 requires AutoEP-managed MoE layers and does not support native DeepSpeed MoE layers, tensor model parallelism from mpu, sequence parallelism, hpZeRO secondary tensor groups, non-1 expert tensor parallelism, or quantized gradients. AutoTP folding (tensor_parallel.autotp_size > 1) works at Stage 3 as at stages 0 to 2, including models built under deepspeed.zero.Init. Stage 3 AutoEP checkpoints are saved partition-natively in the zero_pp_rank_* shard files and support same-topology load, module-only loads (load_module_only), optimizer-state-free loads (load_optimizer_states=False), and Universal Checkpoint conversion, except with AutoTP folding. Optimizer-including Universal Checkpoint loads can resume with a different data-parallel world size, a different autoep_size, or both, when the target autoep_size divides the model’s expert count. Weights-only/module-only Universal Checkpoint loads use the converted fp32.pt parameter files and support the same data-parallel and autoep_size topology changes.

Usage:

{
    "expert_parallel": {
        "enabled": true,
        "autoep_size": 4,
        "preset_model": "mixtral"
    }
}

Experimental regional torch.compile

AutoEP can keep its router, token movement, expert computation, and collectives in eager mode while compiling the surrounding decoder blocks with vanilla torch.compile. This targets fragmented attention, normalization, residual, and dense backward work without capturing AutoEP communication in the graph. The path is opt-in. Enable it in the DeepSpeed configuration, then call engine.compile() after initialization:

{
  "compile": {
    "autoep_non_moe": true
  }
}

engine, optimizer, _, _ = deepspeed.initialize(
    model=model,
    model_parameters=model.parameters(),
    config=ds_config,
)
engine.compile()

The call must happen after deepspeed.initialize() so AutoEP replacement is complete. DeepSpeed discovers each AutoEPMoELayer and regionally compiles its direct callable parent with fullgraph=False and dynamic=False. This is usually a decoder block; a callable model root with a direct AutoEP child is also supported. The AutoEP layer is an explicit compiler-disabled graph break, so routing, AllToAll dispatch/combine, and expert execution remain eager. An AutoEP layer used as the model root has no surrounding region and is rejected.

The decoder blocks contain repeated attention, normalization, residual, and dense work targeted by this optimization. Compiling these regions bounds the traced code and lets structurally identical blocks reuse compiled graphs. Modules outside the selected regions, such as top-level embeddings and the language-model head, remain eager. They are not intrinsically incompatible with compilation, but extending the region would need separate validation and performance measurements. If the selected region is the model root, its non-MoE operations are included.

compile.autoep_non_moe defaults to false, preserving the existing full-model behavior of engine.compile(). Setting the option alone does not compile the model; execution stays eager until engine.compile() is called.

The initial experimental path supports vanilla torch.compile with the standard comm backend, sequence and pipeline parallel sizes of one, and ZeRO stages 0, 1, and 2. Distributed performance and parity validation currently target ZeRO stage 1. It rejects DeepEP, DeepCompile, AutoEP+AutoTP folding, sequence or pipeline parallelism, ZeRO stage 3, optimizer or parameter offload, compiled autograd, DeepCompile schedules, and any fullgraph or dynamic value other than False instead of silently changing the requested behavior.

How it works:

  1. During deepspeed.initialize(), AutoEP scans the model for MoE layers using preset-defined patterns (router name, expert name, weight shapes).

  2. Detected MoE blocks are replaced with AutoEPMoELayer, which uses TorchTitan’s grouped GEMM kernels and AllToAll token dispatch.

  3. EP/EDP process groups are created automatically based on autoep_size.

  4. Expert parameters are marked for expert-data-parallel gradient reduction; router and shared-expert parameters use standard data-parallel reduction.

Router outputs and activation checkpointing:

Models using Hugging Face’s model-level router-logit recording capture the existing gate output; AutoEP does not compute a second projection just to populate an unused cache. Models whose MoE blocks return router logits directly compute them locally when constructing the return value, preserving that return contract and its gradients. No router-logit tensor is stored on the layer, so checkpoint replay early-stop and exceptions cannot leave a router-logit cache keeping the autograd graph alive between training steps.

Communication backend (optional):

The expert AllToAll can be carried by DeepEP instead of the default collectives. This is opt-in and off by default; jobs that set nothing keep the existing path unchanged.

{
  "expert_parallel": {
    "enabled": true,
    "autoep_size": 8,
    "comm_backend": "deepep",
    "comm_num_sm": 12,
    "comm_qp_margin": 4,
    "comm_max_tokens_per_rank": 4096
  }
}

  • 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. Required when comm_backend is "deepep" because the DeepEP buffer is sized statically and must use the same capacity on every rank. A batch that exceeds it is an error.

For autoep_size > 1, DeepEP receives the router output directly, bypassing the collective backend’s sorting, token expansion, and split-count exchange. Shared experts and router-logit outputs retain the same behavior. The EP communicator is initialized once before each layer’s first DeepEP buffer is constructed, including when the caller supplied a lazily initialized process group. This initialization does not run on subsequent forwards. The standard comm and autoep_size=1 paths are unchanged.

Received rows and routing weights are passed through without a prefix slice when their buffers already contain exactly the arrived rows. Overallocated buffers are still trimmed independently to the handle’s row count. This avoids redundant slice-gradient zero-fills and copies during checkpoint recomputation; it changes neither row weighting nor DeepEP’s communication volume.

Both reentrant and non-reentrant activation checkpointing are supported with DeepEP. The order in which DeepEP delivers rows depends on arrival timing, and a non-reentrant checkpoint keeps the forward pass’s autograd nodes, together with the dispatch handle they hold. A non-reentrant recompute therefore repeats the forward’s dispatch on that handle rather than dispatching afresh, so the expert activations it restores are in the order the backward expects. A recompute is matched to its forward pass by requiring each micro-batch to run backward before the layer’s next forward, as gradient accumulation does; a recompute that could belong to more than one live forward pass raises an error rather than guessing. This also rejects interleaving a retained graph with a newer forward, even if the retained graph has already been replayed once. Repeating backward on one retained graph without another live forward remains supported. An inner non-reentrant checkpoint nested under a reentrant checkpoint records its fresh graph-building dispatch during the outer backward and reuses that layout for the inner recompute. Use reentrant checkpointing for schedules that run several forwards before their backwards. With PyTorch’s default early stop, a non-reentrant recompute ends once the expert activations are restored, so it does not repeat the combine when nothing after the combine in the checkpointed region saves activations, for example a layer without shared experts.

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.

Within one model, all MoE layers that agree on EP group, expert count, top-k, hidden size, capacity, comm_num_sm and comm_qp_margin share a single DeepEP buffer, which is every layer of a normal model. A buffer reserves fabric resources that are not reported as device memory and that run out: measured on 32 H100s across four nodes, the twenty-eighth buffer per rank fails inside ncclDevCommCreate, so one buffer per layer put a 27-layer ceiling on the backend there. Buffers are also slow to build, about 15 seconds each on 16 H100s and 22 on 32, so sharing removes minutes of startup as well. The buffer is released once the last layer holding it is torn down.

Sharing stops at the model. Two models converted separately get their own buffers even on the same EP group with identical geometry, because they are driven independently: an actor and a frozen reference model in a reinforcement learning loop need not reach their MoE layers in any fixed order relative to each other, and a shared DeepEP communication context would make that order matter. The cost is one extra buffer per model against a ceiling of 27.

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.

Fused gate and up projections (experimental):

The grouped experts compute two projections of the same input, gate with w1 and up with w3, before the expert activation. gate_up_impl selects whether they run as two grouped GEMMs or as one:

{
  "expert_parallel": {
    "enabled": true,
    "autoep_size": 8,
    "gate_up_impl": "fused"
  }
}

"separate" (default) keeps the existing two GEMMs. "fused" concatenates w1 and w3 on each call and runs one torch._grouped_mm with twice the output columns. The activation then reads the two halves of that output in place: swiglu uses a Triton kernel, and the other forms run their usual expression on views of the halves. The parameters are not changed, so checkpoints, universal checkpoints, expert tensor parallelism and optimizer state, including Muon’s per-expert treatment of w1 and w3, see the same layout as before.

The forward expression and parameter-gradient semantics are unchanged. Exact output and weight-gradient equality was observed in the tested H100 cases, not guaranteed for all GEMM backends, shapes or hardware. In the backward pass, one input-gradient GEMM accumulates both projections in FP32, where "separate" adds two results rounded to the input dtype, so the input gradient can differ in the last bits. The numerical regression compares relative L2 errors with an FP64 reference using the same represented BF16 parameters; it does not claim universal bitwise equivalence or accuracy improvement. Forward and parameter-gradient checks use the grouped-GEMM tests’ existing BF16 bound (2% relative plus 2% of the reference tensor’s maximum magnitude). The clamped activations’ parameter gradients use an independent dense BF16 reference because rounding a projection across the clamp boundary changes its derivative discontinuously; smooth forms use the FP64 reference. The input-gradient relative-L2 bound is unchanged. Cross-algorithm training checks allow the same rounding scale, while reloading each parameter and optimizer momentum tensor must still preserve its saved values exactly.

The packed SwiGLU Triton path supports first-order gradients only: create_graph=True raises instead of silently dropping its second derivative. Forward and backward launch on the input device’s current stream, restoring the caller’s device. Other activations keep their existing eager expression.

Most of the saving is in the backward pass: one input-gradient GEMM instead of two, and no separate addition. On Qwen3-30B-A3B with 8 H100s (expert parallel size 8, micro-batch 4, 16 accumulation steps, sequence length 4096, DeepEP and reentrant activation checkpointing), a full optimizer step went from 22.02 s to 21.60 s (-1.9%) in two pairs using an equivalent benchmark prototype in the same allocation, with unchanged peak allocated memory and about 0.4 GB more reserved memory.

"fused" is rejected, rather than silently ignored, when the experts would not run torch._grouped_mm: with use_grouped_mm=false, or on devices that select the Triton grouped GEMM (compute capability below 9.0) unless disable_triton_grouped_mm=true.

DeepEP row weighting implementation (experimental):

DeepEP dispatch returns one received row per routed assignment and one FP32 weight per row. row_weighting_impl selects how AutoEP multiplies those rows by their weights at the existing score_apply boundary:

{
  "expert_parallel": {
    "enabled": true,
    "autoep_size": 8,
    "comm_backend": "deepep",
    "comm_max_tokens_per_rank": 4096,
    "row_weighting_impl": "fused"
  }
}

"auto" (default) resolves to "eager", preserving the existing eager expression exactly. "fused" runs a separate Triton pointwise operator for (rows.float() * weights).to(rows.dtype). It does not reduce over top-k, does not change where BF16/FP16 rounding occurs, and does not replace DeepEP’s combine; the output remains one weighted row per received row in the same row order.

The forward product and row gradient match eager’s rounding. The FP32 gradient of the routing weight sums the same products in a different order, so it need not be bitwise equal to eager’s; neither summation is consistently closer to an FP64 reference. Comparisons should use gradient errors relative to the gradient norm after backward and before optimizer clipping in engine.step(). Adam’s first update can differ on the scale of the learning rate when a near-zero gradient changes sign, even if the overall gradients agree closely.

"fused" is rejected, rather than silently ignored, when AutoEP cannot honor it:

  • comm_backend is not "deepep" or autoep_size=1, because the call sites exist only inside the DeepEP route;

  • Triton is unavailable, the device is not CUDA, or the build is ROCm;

  • rows are not bfloat16 or float16;

  • weights are not FP32 [N, 1] tensors on the same CUDA device;

  • rows or weights are not contiguous, or rows are not shaped [N, H].

The operator also accepts FP16 rows, but the current DeepEP dispatch supports BF16 rows only. Correct backward replay through DeepEP additionally requires preserving the cached dispatch layout; that correction is independent of row weighting. The separate MoE gradient-norm correction affects the FP16_Optimizer wrapper, which is also used by some BF16 configurations (for example, BF16 with BF16 gradient accumulation without ZeRO). The model-level gradient comparison samples gradients before the wrapper computes the norm and clips them in engine.step(). GPU validation applies both independent corrections; neither is part of this opt-in change.

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 takes bfloat16 rows, so selecting this backend for an fp16 or fp32 run is rejected rather than silently downgraded.

  • Not compatible with folded tensor parallelism (expert_tensor_parallel_size > 1), which is rejected at setup.

Python cyclic GC (experimental):

Large Python model graphs can accumulate cyclic objects during training, and a generation-2 collection pauses one rank’s Python thread, which is then exposed as collective wait time on every expert-parallel rank. The top-level disable_python_gc option addresses this. It is process-wide rather than AutoEP-specific, so it is documented with the general configuration options.

Fused RMSNorm (experimental):

RMSNorm is a model module rather than an AutoEP expert-parallel component, so it is not configured under expert_parallel. DeepSpeed provides an opt-in installer that runs Hugging Face RMSNorm modules with fused Triton kernels:

from deepspeed.ops.triton_ops.fused_rms_norm import replace_rms_norm

replaced = replace_rms_norm(model)

The installer adds a process-wide class dispatcher, with a per-instance opt-in marker. Other instances of those classes keep their original eager forward until explicitly selected. Shallow/deep copies and torch.nn.DataParallel replicas retain their live opt-in and use their own weights and epsilon rather than the original receiver. The HF module classes, Parameters and state-dict keys are unchanged. Full-model checkpoints saved with torch.save(model, ...) can be loaded with torch.load(..., weights_only=False). Loading clears the opt-in, so the loaded model runs eager until replace_rms_norm(model) is called again, even if this process already has a dispatcher installed for another instance. No fused dispatcher needs to be serialized.

The kernels reproduce the RMSNorm expression of Llama-style models: the variance and normalization are computed in FP32, the normalized value is cast back to the input dtype once, and only then multiplied by the weight (gamma). A class name ending in RMSNorm with a 1-D weight and a variance_epsilon does not imply that order: GptOssRMSNorm, for one, multiplies by the weight before the cast, and Olmo2RMSNorm moved to that order between Transformers releases. The installer therefore replaces only instances of the classes in SUPPORTED_RMS_NORM_CLASSES: the RMSNorm classes of Llama, Mistral, Mixtral, Phi-3, Qwen2, Qwen2-MoE, Qwen3, Qwen3-MoE, DeepSeek-V2 and DeepSeek-V3. A module is replaced only if it still runs its class’s own forward and is at most 2048 wide, so subclasses, modules whose forward was already patched or wrapped (by another kernel installer or by an earlier call) and wider norms are left alone. Replaced modules keep their own weight Parameter and variance_epsilon. The return value is the number of modules replaced: 193 for Qwen3-30B-A3B, four per decoder layer plus the final norm.

A replaced module runs the kernels when its input and weight are bfloat16 or float16 CUDA tensors of the same dtype and Triton is available. For any other input it runs its class’s own eager forward, so the result is exactly eager’s, and logs a warning once. The installer can therefore run before the model is moved to the GPU. fused_rms_norm(hidden, weight, eps) in the same module applies the kernels directly and raises on unsupported inputs instead.

Where the kernels run, results are not bitwise identical to eager RMSNorm: the kernels’ FP32 reductions sum in a different order, so a small fraction of outputs and gradients differ in their last bits.

Requirements and limits:

  • The kernels need CUDA with Triton. On ROCm or without Triton, replaced modules run eager.

  • bfloat16 or float16, with the input and the weight in the same dtype. Other inputs run eager, including float32 weights under autocast.

  • Norms wider than 2048 are not replaced. That covers the hidden size and head dimension of Qwen3-30B-A3B; in models with a wider hidden size, such as Qwen3-235B-A22B (4096), only the narrower norms, such as the 128-wide query and key norms, are replaced.

  • Any leading shape. Non-contiguous inputs and weights are copied before the kernels run. The output and the input gradient the kernels produce are contiguous, whereas eager RMSNorm follows the strides of its input and of the incoming gradient.

  • First-order gradients for the input and the weight; the epsilon is a constant. Differentiating those gradients again (double backward) raises. torch.compile and torch.func transforms are not covered.

Fused weighted restore (experimental):

After the combine all-to-all, AutoEP holds one row per routed assignment and has to turn it back into one row per token. combine_impl selects how:

{
    "expert_parallel": {
        "enabled": true,
        "autoep_size": 16,
        "preset_model": "qwen3_moe",
        "combine_impl": "fused_weighted_sum"
    }
}

"auto" (default) resolves to "weighted_sum", which scatters the rows into a zero-filled [tokens * top_k, hidden] buffer, widens it to FP32 to apply the routing weights, and reduces over top-k. "fused_weighted_sum" computes the weighted reduction in one kernel launch: each program owns a whole token and walks the hidden dimension in chunks, multiplying each routed row by its score in FP32 and accumulating the products in slot order in FP32. FP contraction is disabled so each product rounds before it is added; the result is cast only once to the output dtype. Neither the scattered buffer nor the FP32 intermediate is allocated. At the canonical shape the FP32 intermediate alone is 64 MiB per layer.

The top-k summation order can differ from eager, so results are within the existing numerical tolerances, not bitwise unchanged. Only the forward reduction changes: the backward, collectives, router, grouped GEMM and expert-major reorder are untouched.

"fused_weighted_sum" is rejected, rather than quietly ignored, when it would have nothing to replace or would change semantics:

  • tensor_parallel.autotp_size greater than 1, which uses folded tensor parallelism and restores combined tokens from assignment metadata instead;

  • expert_tensor_parallel_size greater than 1;

  • comm_backend="deepep" with expert parallelism, because DeepEP already restores and reduces its routed rows;

  • a resolved score_apply other than "post";

  • activations that are not bfloat16, float16, or float32, a non-CUDA device, or a build without Triton.

Failing fast matters for measurement: a run that asked for the fused reduction and silently got the eager one would report the difference between an implementation and itself.

Fused rotary position embedding (experimental):

RoPE belongs to the attention layers rather than to AutoEP, so it is not configured under expert_parallel. DeepSpeed provides an opt-in installer that runs Hugging Face’s apply_rotary_pos_emb with a fused Triton kernel:

from deepspeed.ops.triton_ops.fused_rotary_pos_emb import replace_rotary_pos_emb

patched = replace_rotary_pos_emb(model)

The eager function computes q * cos + rotate_half(q) * sin in several elementwise kernels, each reading and writing the whole query or key tensor. The fused kernel reads each tensor once and writes it once, in the forward and in the backward, and applies each position’s cos and sin to all of its heads. Every product and sum is rounded to the input dtype where the eager expression rounds it, so the outputs, and the gradients for the queries and keys, equal eager’s element for element as compared by torch.equal, which does not distinguish +0.0 from -0.0.

Attention modules look apply_rotary_pos_emb up in their modeling module, so the installer replaces it there: the replacement applies to every model of that architecture in the process, not only to model, and restore_rotary_pos_emb() undoes it. A modeling module is patched only if one of the model’s submodules is defined in it, it is listed in SUPPORTED_ROTARY_MODULES (Llama, Mistral, Mixtral, Qwen2, Qwen2-MoE, Qwen3, Qwen3-MoE and DeepSeek-V3), and its apply_rotary_pos_emb and rotate_half still have the code of the split-half expression: the same bytecode, names, constants and defaults, whatever their docstrings. Other architectures define functions of the same name that rotate only part of the head dimension or add casts, so a module whose function has changed is left alone with a warning. The return value is the number of modeling modules patched: 1 for Qwen3-30B-A3B.

The replacement runs the kernel when the queries, keys, cos and sin are bfloat16 or float16 CUDA tensors of one dtype, the head dimension is even, at most 512 and has unit stride, and cos and sin do not require grad. For any other input it runs the original function, so the result is exactly eager’s, and logs a warning once. fused_apply_rotary_pos_emb in the same module applies the kernel directly and raises on unsupported inputs instead.

Requirements and limits:

  • The kernel needs CUDA with Triton. On ROCm or without Triton, the replaced function runs eager.

  • All four tensors must be on one device. Launches use that device’s current stream even if another device is current, and restore the caller’s current device afterwards.

  • unsqueeze_dim 1 (heads before the sequence) or 2 (sequence before the heads), with cos and sin of shape [batch, sequence, head_dim] or [1, sequence, head_dim].

  • For dense layouts, such as attention’s transposed projection, outputs keep the strides of the queries and keys, as eager’s do; other non-contiguous inputs can give outputs laid out differently from eager’s. The gradients for the queries and keys take the layout of the outputs, whereas eager’s follow the incoming gradient, so weight gradients computed from them further back, such as the query projection’s, can differ from eager’s in the last bits.

  • Transformers releases whose apply_rotary_pos_emb still takes position_ids, such as 4.51, are left unpatched.

  • First-order gradients for the queries and keys; cos and sin are constants. Differentiating those gradients again (double backward) raises. torch.compile and torch.func transforms are not covered.

Constraints:

  • autoep_size must divide num_experts for all detected MoE layers.

  • autoep_size=1 is valid: all experts remain local (no AllToAll), useful for functional testing on a single GPU.

  • AutoEP combines with AutoTP (tensor_parallel.autotp_size > 1) by folding: on the same ranks, AutoTP shards the layers its partition specification names and AutoEP distributes the routed experts. Folding supports ZeRO stages 0 to 3 with expert_tensor_parallel_size=1, and rejects ZeRO optimizer or parameter offload and DeepCompile. Tensor model parallelism from mpu is not supported.

  • AutoEP with ZeRO Stage 3 is supported only without sequence parallelism, hpZeRO secondary tensor groups, non-1 expert tensor parallelism, or quantized gradients.

  • Regular checkpoint save/load requires matching autoep_size. To change autoep_size or data-parallel world size across runs for the same AutoEP-detected model topology, convert the checkpoint to Universal Checkpoint format and load it with checkpoint.load_universal; see the Universal Checkpointing tutorial for the detailed flow and constraints.

  • DeepSeek-V2 and DeepSeek-V3 AutoEP do not support load-balance expert bias yet. The built-in DeepSeek presets disable it by default; explicit non-null values fail.

Chunked causal-LM loss

Long sequences with large vocabularies spend much of their peak memory on the language-model loss rather than on the model. Hugging Face’s ForCausalLMLoss upcasts the full [tokens, vocab] logits to FP32 and keeps an FP32 log-softmax for backward, and backward allocates two more FP32 tensors of the same size. At 8,192 tokens and Qwen3’s 151,936-token vocabulary each is about 5 GB, and all of them are live when backward starts, which is where AutoEP training memory peaks.

install_chunked_causal_lm_loss replaces that loss with one that computes the same FP32 cross entropy without them: forward keeps only the logits it was given plus one FP32 value per token, and backward recomputes the softmax and writes the gradient directly in the logits’ dtype. On CUDA a Triton kernel does each pass in a single read of the logits (backend="triton"); the PyTorch backend (backend="torch") does the same arithmetic in blocks of rows and runs anywhere. The default "auto" picks Triton when it can run; naming a backend that cannot run raises. The loss, the label shifting, ignore_index and num_items_in_batch normalization follow ForCausalLMLoss; gradients may differ from it by one BF16/FP16 rounding step where the FP32 evaluation order lands on a rounding boundary. It has no second derivative: backward with create_graph=True raises. It is opt-in, and it refuses models whose loss_function is not the stock ForCausalLMLoss:

from deepspeed.runtime.chunked_cross_entropy import install_chunked_causal_lm_loss

install_chunked_causal_lm_loss(model)
engine, optimizer, _, _ = deepspeed.initialize(model=model, config=ds_config)

Class-index targets must have dtype torch.long. Other dtypes, including floating-point and boolean labels, are rejected instead of being converted. When shift_labels is provided, it is the target used by the loss; unused labels are not dtype-checked. block_rows must be a positive integer or None for automatic sizing, even when the selected backend does not use row blocking. Invalid settings are rejected before a backend runs. A model with a non-writable stock loss_function raises an actionable error rather than being silently left unmodified.

Planning expert placement transfers

An AutoEPMoELayer can describe where its live expert weights are, and plan the copies that would move them to another expert placement, for example one with a different EP size. Planning moves no data, does not plan optimizer states, and does not change the placement the layer runs with.

from deepspeed.checkpoint.autoep_affine import make_autoep_placement_descriptor

# layer: an AutoEPMoELayer after deepspeed.initialize()
affine_maps = layer.expert_affine_maps()  # {"w1": ParamAffineMap, "w2": ..., "w3": ...}

# Move 8 experts onto 2 EP ranks: rank 0 holds experts 0-2, rank 1 holds experts 3-7.
target = make_autoep_placement_descriptor(8, [[0, 1, 2], [3, 4, 5, 6, 7]])
plans = layer.plan_expert_transfer(target)  # {"w1": TransferPlan, "w2": ..., "w3": ...}
for segment in plans["w1"].segments:
    ...  # copy segment.shape elements from the source shard into the target shard

  • make_autoep_placement_descriptor(num_experts, experts_by_rank) builds the target placement. experts_by_rank[r] lists the global expert IDs that target rank r holds, in the order they are packed in its local tensor. Ranks may hold different numbers of experts, non-contiguous experts, no experts, or experts that other ranks also hold, but every expert must be held by at least one rank.

  • Ranks in a plan are EP-local. segment.source_rank is a rank of the layer’s EP group; translate it with dist.get_global_rank(layer.ep_group, segment.source_rank). segment.target_rank indexes experts_by_rank of the target placement; mapping it to a process is up to the caller, since the target EP group may differ from the current one.

  • source_offset/target_offset and the stride vectors count elements of a rank’s packed local expert tensor [local_experts, ...] in row-major order. That tensor is the unpartitioned parameter: under ZeRO-3, gather the parameter (for example with deepspeed.zero.GatheredParameters) before applying the offsets. Ranks are EP-local, so expert-data-parallel replicas are not candidates here; plan from expert_storage_affine_maps() to read from them.

  • When several source ranks hold an element, the plan reads it from the lowest-ranked holder. Pass replica_selector, a callable that receives a ReplicaSelectionRequest and returns the source_rank of one of its candidates, to choose a different holder.

  • limits takes a deepspeed.checkpoint.affine_transfer.PlanBudget. A plan that exceeds it raises AffineTransferError instead of being truncated.

  • The runtime packs experts with the contiguous, uniform placement. If the layer’s placement descriptor describes anything else, both methods raise NotImplementedError.

expert_affine_maps() describes one packed tensor per EP rank, so ZeRO and expert-data-parallel (EDP) storage has to be normalized before its offsets apply. expert_storage_affine_maps() instead describes what every process actually stores, so a plan built from it reads directly from runtime storage:

from deepspeed.checkpoint.affine_transfer import plan_transfer

# Built locally from the expert data parallel group table; no communication.
storage_maps = layer.expert_storage_affine_maps()  # {"w1": ParamAffineMap, ...}
plan = plan_transfer(target_map, storage_maps["w1"])

  • Map ranks are global ranks. The logical shape is the parameter’s own shape, and each shard is a rank’s flat storage, which zero_param_storage_shard(param) in deepspeed.checkpoint.zero_affine returns: param.data.view(-1), or under ZeRO-3 the partition param.ds_tensor without its padding. target_map must use the same logical shape; autoep_storage_affine_map in deepspeed.checkpoint.autoep_affine builds one from a placement descriptor and its holders.

  • Under ZeRO-0/1/2, every EDP replica of an EP rank stores the whole packed tensor and is listed as a holder. The plan reads from the lowest global rank unless replica_selector chooses another candidate.

  • Under ZeRO-3, each EDP rank stores its partition of the packed tensor. The trailing padding is not part of a shard, so a rank whose partition is all padding, or whose EP rank holds no experts, has an empty shard and is never a source or target of a copy.

  • The maps list only authoritative copies: the storage the optimizer step updates and the next forward reads, so a write through the maps lasts. Copies the runtime derives from them are not listed. One is the gathered copy ZeRO-3 keeps of a persistent parameter between forwards: it holds the same values after a step, but a write to it is lost once it is released, and a write to the shard alone is not seen while it serves the forward. For this reason zero_param_storage_shard refuses a gathered parameter; call engine.empty_partition_cache() before writing.

  • The maps describe the low-precision parameters the forward reads, which is what a rollout engine consumes. When a training engine is the target, these parameters are themselves derived from its fp32 master weights, and its next optimizer step overwrites them. The fp32 master weights and optimizer states are not described. Expert tensor parallelism (expert_tensor_parallel_size > 1) raises NotImplementedError.

  • autotp_storage_affine_maps(model) in deepspeed.checkpoint.zero_affine describes the experts of an AutoEP layer with this method, so a model-wide description under AutoTP and AutoEP is consistent.

The maps and plans use the affine representation that Universal Checkpoint uses for expert parameters; see deepspeed/checkpoint/affine_ir_spec.md.