A mixture-of-experts layer breaks the assumption every other parallelism strategy is built on: that all ranks hold the same parameters and do the same work on different data. Expert parallelism puts different experts on different GPUs, so the feed-forward block a token needs may not live on the device that holds the token. On the serving side that shows up as a decode-latency problem. On the training side it shows up somewhere less obvious: in the shape of the device mesh, in the fact that your model now has two classes of parameter that synchronize over two different process groups, in optimizer state that shards differently for experts than for everything else, and in a step time set by whichever rank got the most tokens. This piece is about that training step — how the mesh is laid out, what synchronizes with whom, and where the stalls come from.
The mesh: EP, TP and DP are three axes of one device grid
Think of the cluster not as a list of GPUs but as a grid with one axis per parallelism strategy. Tensor parallelism splits a single matrix multiply across a few devices, so a TP group is small and lives inside one node where NVLink bandwidth is cheap. Expert parallelism splits the expert pool across devices: with 64 experts and an EP degree of 8, each rank in the EP group owns 8 of them. Data parallelism replicates the whole arrangement over distinct microbatches.
These are not alternatives — they multiply. A 512-GPU job might run TP = 8, EP = 8, DP = 8, and every device belongs simultaneously to a TP group, an EP group and a DP group. Pipeline parallelism adds a fourth axis by cutting the model into stages; that dimension is covered in gpu_pipeline_parallel. The layout question is not which strategy to use but which axis gets mapped onto which physical link, because each axis carries a different collective with a different tolerance for latency.
Where one device sits in all three groups
Rank placement is the decision that determines your communication bill. The ordering that usually wins makes TP the fastest-varying dimension so a tensor-parallel group lands entirely within a node and its per-layer reductions never touch the network. EP is usually placed next, spanning nodes if it must, because its all-to-all is bandwidth-heavy but happens twice per MoE layer rather than twice per linear. DP varies slowest, so the data-parallel all-reduce is the collective that crosses the most switch hops — acceptable, because it fires once per step and can be overlapped with backward compute.
A single GPU therefore holds a TP shard of a shared weight, a subset of the expert pool, and one replica's worth of activations, all at once. Get the ordering backwards — EP inside the node, TP across nodes — and you have moved the most frequent collective onto the slowest fabric. See gpu_nccl_collectives for how the underlying algorithms behave on each topology.
Sharded by expert, replicated everywhere else
An MoE transformer has two parameter populations, and almost every subtlety in this article follows from the split. Expert parameters — the per-expert feed-forward weights, which typically dominate the total parameter count — are partitioned by expert across the EP group. Rank 0 physically does not have expert 17's weights and never will. Shared parameters — attention projections, embeddings, layer norms, and the router itself — are replicated across the EP group exactly as in a dense model, and sharded only along the TP axis where applicable.
The router deserves a note: it is tiny, it is shared, and it must produce identical decisions on every rank in the EP group, because those ranks are all dispatching from the same logical token batch. Its gradient behaves like any other shared parameter. The routing rule that produces the assignment is the subject of gpu_moe_routing and, mathematically, of the transformer_math treatment of Switch-style routing.
Two gradient reductions over two different groups
This is the point that most often gets implemented wrong. Shared parameters are replicated across the full data-parallel width, so their gradients all-reduce over the whole DP group, just as in dense training. Expert parameters are not replicated that widely. A given expert's weights exist only on the ranks that were assigned that expert, so its gradient must be reduced over precisely that set — the expert data-parallel group, whose size is roughly world size divided by (EP × TP), not world size divided by TP.
Two consequences follow. First, the expert gradient reduction moves far less data than a naive full-width all-reduce would, which is part of why EP is affordable at all. Second, if EP is set wide enough that every expert lives on exactly one rank per replica group, that group has size one and the expert gradients need no cross-rank reduction whatsoever. Collapsing both populations onto a single DP process group silently averages weights that were never supposed to be averaged.
Optimizer state follows the expert group, not the world
Adam-style optimizers carry roughly twice the parameter count in moments plus, commonly, a high-precision master copy — state that dwarfs the weights themselves. Sharding it is mandatory at MoE scale, and the sharding must follow the same group structure as the gradients. Shared-parameter optimizer state shards across the full data-parallel group. Expert optimizer state shards across the expert data-parallel group only, because that is the set of ranks that holds a copy of those weights in the first place.
The practical effect is that the memory win from partitioning optimizer state is much smaller for experts than for shared parameters, since the group you can partition across is smaller by a factor of EP. When EP already places each expert on a single rank, there is nothing left to partition and the state sits whole. gpu_zero_sharding covers the partitioning stages themselves; the MoE wrinkle is purely which group each stage applies over.
The step timeline: dispatch, expert GEMM, combine
Inside one MoE layer's forward pass the sequence is fixed. The router scores every token; tokens are permuted into contiguous per-destination buffers; an all-to-all sends each token to the rank owning its expert; each rank runs its local experts over whatever arrived; a second all-to-all returns the outputs; the inverse permutation restores the original token order. Backward walks the same path in reverse, so a training step pays four all-to-alls per MoE layer, not two.
The permutation step matters more than it looks. Sorting tokens by destination is what lets the expert computation become a batched or grouped matrix multiply instead of many ragged small ones — the kernel side of that is gpu_moe_grouped_gemm. The collective itself, and its cost model, belong to gpu_moe_all_to_all.
Overlap or contention: one fabric, two collectives
Standard data-parallel training hides the gradient all-reduce behind backward compute: as each layer's gradients become ready they are bucketed and reduced while later layers are still computing. That trick still works for shared parameters in an MoE model, but it now competes with something. The backward pass of every MoE layer is itself issuing all-to-alls, and both collectives want the same NICs and the same links.
So “overlapping communication with computation” quietly becomes “overlapping communication with communication.” If the shared-parameter reduction is already saturating the interconnect, an all-to-all issued into the same window does not hide — it queues, and both finish later than either would alone. The usual mitigations are to schedule the two onto separate streams with explicit priorities, to delay or chunk the gradient reduction so it fills gaps between dispatches rather than colliding with them, and to keep the EP axis on links the DP reduction is not already consuming.
Recomputation re-runs the dispatch unless you save the permutation
Activation checkpointing trades compute for memory by discarding a block's intermediate activations and recomputing them in backward. Applied naively to an MoE block, it recomputes the whole block — including the router, the permutation and the dispatch all-to-all. That converts a memory saving into a communication cost, because the token shuffle is replayed rather than replaced, and an already communication-bound step gets another round of all-to-all traffic per checkpointed layer.
The fix is selective. Treat the routing decision as cheap-to-store and expensive-to-reproduce: keep the assignment indices and the permutation map, which are small integer tensors, and discard only the large hidden-state activations inside the expert MLPs. Backward then reuses the saved indices and recomputes only local arithmetic. There is a correctness argument for this too — any routing that depends on batch statistics or on a stochastic tie-break is not guaranteed to reproduce the same assignment on a second evaluation, so saving the map removes a class of silent mismatch as well as the traffic.
The straggler rank that sets step time
Every collective in the step is a barrier. The combine all-to-all cannot complete until every rank has finished its local expert GEMM, and the shared-parameter all-reduce cannot retire until every rank has cleared the MoE layers. So the duration of a step is set by the slowest rank, and under expert parallelism the slowest rank is whichever one happened to receive the most tokens this microbatch.
Routing is data-dependent, so that skew is not fixed — it moves batch to batch and drifts as the model trains. The damage is superlinear in a way that is easy to miss: a rank at twice the average token count does not cost you a small percentage, it costs every other rank in the group an idle wait of the same length, repeated at every MoE layer. Profilers show this as a wide, ragged gap before each combine that looks like slow communication but is really compute waiting on one peer. The mechanisms that flatten the distribution — the balancing loss the diagram names, capacity limits, and drop-free variants — are covered in gpu_moe_load_balance and gpu_dropless_moe.