Expert Parallelism for Distributed Training

While standard data parallelism effectively scales training across multiple devices by processing different data batches on each, it requires that a complete copy of the model resides on every device. For a Mixture of Experts model with hundreds of billions or even trillions of parameters, this is a non-starter; the memory requirements of the full parameter set far exceed the capacity of any single accelerator.

This fundamental memory constraint necessitates a different approach. Expert parallelism is a specialized distributed training technique designed specifically to address this challenge. Instead of replicating the entire model, expert parallelism shards the experts across the available devices. Each device holds only a fraction of the total experts, while the non-MoE parts of the model (like the self-attention blocks and the gating network) remain replicated.

The All-to-All Communication Pattern

The defining characteristic of expert parallelism is its reliance on a communication collective known as All-to-All. When a token needs to be processed by an expert that resides on a different device, it must be physically sent over the network. This happens in a two-stage process within every MoE layer.

Token Dispatch (First All-to-All): After the replicated gating network on each device determines the expert assignment for its local batch of tokens, a global data shuffle occurs. Each device sends its tokens to the specific devices that house their assigned experts.
Result Aggregation (Second All-to-All): Once the remote experts have processed their assigned tokens, the resulting output representations are sent back to their original source devices.

This process ensures that each token is processed by the correct expert, regardless of where that expert is located, and that the token representations are returned to their original sequence order to continue the forward pass.

The diagram below illustrates this data flow across two devices (GPUs), each holding a distinct set of experts.

Data flow in expert parallelism. Local tokens are routed to experts on other devices via a first All-to-All communication step (red). After processing, the results are returned to their original devices via a second All-to-All step (blue).

Performance Implications and Trade-offs

Expert parallelism elegantly solves the memory problem, allowing models to scale to trillions of parameters. However, this benefit comes at the cost of significant communication overhead. The All-to-All operation is one of the most bandwidth-intensive collectives in distributed computing because every device must communicate with every other device simultaneously.

The performance of this operation is bottlenecked by the interconnect bandwidth between devices (e.g., NVLink for GPUs within a node, or the network fabric for multi-node setups). The total data volume being shuffled is proportional to the number of tokens, their hidden dimension, and the number of experts each token is routed to ( $k$ in Top-k routing).

This introduces a critical trade-off:

Memory Savings: You can store larger models.
Communication Cost: Your training throughput is now sensitive to network bandwidth, not just computational (FLOPs) capacity.

In practice, this means that simply adding more devices does not guarantee a linear speedup. The efficiency of the All-to-All shuffle becomes a dominant factor in overall training time.

Implementation in Code

While the underlying logic is complex, modern deep learning frameworks designed for large-scale training abstract away much of the difficulty. Libraries like DeepSpeed or fairscale provide modules that handle the expert sharding and communication automatically.

Below is a simplified pseudo-code representation of what happens inside a distributed MoE layer.

# Pseudo-code for an expert-parallel forward pass on a single device

# Non-MoE layers are processed as usual
x = self_attention(input_tokens)

# Gating network is replicated on all devices
# It computes assignments for its local tokens
gate_logits = gating_network(x)
local_assignments = top_k_router(gate_logits)

# First All-to-All: Shuffle tokens to the devices that own the target experts.
# The framework handles the low-level communication.
# Each device receives a batch of tokens intended for its local experts.
shuffled_tokens = all_to_all_dispatch(x, local_assignments)

# Process the received tokens using the local subset of experts.
# These experts only exist on this device.
processed_shuffled_tokens = local_experts(shuffled_tokens)

# Second All-to-All: Shuffle the processed tokens back to their original devices.
# This restores the original token order.
output_tokens = all_to_all_combine(processed_shuffled_tokens, local_assignments)

# The rest of the model proceeds
...

This model-sharding strategy is the foundation for training truly massive MoE models. However, it is rarely used in isolation. To achieve maximum efficiency, expert parallelism is almost always combined with other parallelism techniques, a topic we will cover in the next section.

Was this section helpful?

References

Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer, Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, Jeff Dean, 2017 arXiv preprint arXiv:1701.06538 DOI: 10.48550/arXiv.1701.06538 - Introduces the sparsely-gated Mixture-of-Experts layer, a foundational work for MoE models and their large-scale application.
GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding, Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, Zhifeng Chen, 2020 arXiv preprint arXiv:2006.16668 DOI: 10.48550/arXiv.2006.16668 - Describes a system for automatically sharding large models, including experts, across multiple devices, and details the all-to-all communication pattern essential for expert parallelism.
Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity, William Fedus, Barret Zoph, Noam Shazeer, 2022 Journal of Machine Learning Research, Vol. 23(120) (JMLR) DOI: 10.5555/3545945.3546197 - Presents Switch Transformers, demonstrating the practical scalability of MoE models to trillions of parameters using distributed training strategies, including expert parallelism.
Mixture of Experts (MoE) Layer, DeepSpeed Team, 2024 (Microsoft) - Official DeepSpeed documentation on implementing Mixture of Experts layers, covering the framework's approach to distributed training for MoE models.