Part of Language AI Handbook
Explains how expert parallelism distributes MoE experts across devices using all-to-all communication, enabling efficient training of trillion-parameter models.
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Expert Parallelism
Mixture of Experts models achieve their remarkable efficiency by activating only a subset of parameters for each token. As we explored in previous chapters, a model with 64 experts might only use 2 experts per token, giving you the capacity of a massive model with the compute cost of a much smaller one. A single forward pass through the model touches only a fraction of the total parameter count, yet the model has access to all of the capacity when it needs it. This conditional computation property is what makes MoE architectures so appealing: you pay for what you use, not for everything you have.
But this creates a basic systems challenge that becomes acute at scale. As models grow to hundreds of billions or even trillions of parameters spread across dozens of expert networks, where do all those experts physically reside? No single GPU can hold them all. And how do tokens get to the right experts when those experts might live on a completely different device, possibly in a different server rack across the data center? The clever routing mechanism that makes MoE efficient at the algorithmic level creates a deep systems problem at the hardware level.
Expert parallelism is the distributed-computing strategy designed specifically to solve this problem. It is the engineering foundation that makes large-scale MoE training possible. Unlike data parallelism, which replicates the entire model across devices and processes different data on each copy, or tensor parallelism, which shards individual weight matrices across devices, expert parallelism distributes the expert networks themselves. Each device in an expert-parallel setup becomes responsible for a specific subset of experts. Tokens travel to wherever their selected experts live, computation happens locally on the hosting device, and results travel back. This chapter examines how this works in practice: from assigning experts to devices, through the all-to-all communication patterns that move tokens across the cluster, to the implementation details that determine hardware utilization.
Understanding expert parallelism requires thinking simultaneously at two levels: the algorithmic level, where we reason about routing decisions and expert computations, and the systems level, where we reason about network bandwidth and latency as well as hardware utilization. The two levels are deeply interconnected. A routing decision that looks elegant at the algorithmic level can be a performance disaster at the systems level if it creates lopsided communication traffic. Conversely, systems optimizations like capacity factors and token dropping have direct implications for model quality. Getting expert parallelism right means understanding both levels and the tensions between them.
Think of expert parallelism like a library system distributed across multiple branch locations. Each branch specializes in different subjects, holding the relevant books. When a patron submits a research request (a token), a routing coordinator (the gating network) determines which branches have the most relevant materials (which experts to consult). The patron's request travels to those branches, the relevant books are retrieved, and the results come back to the patron. The coordinator does not need to know the contents of every book; it just needs to know which branches specialize in what. The efficiency of this system depends critically on how well the branches are stocked relative to demand, how quickly requests can travel between locations, and whether popular branches become overwhelmed with requests while others sit idle.
Expert parallelism as a practical training strategy emerged from the Switch Transformer work at Google Brain in 2021, which demonstrated that MoE models with trillions of parameters could be trained efficiently by assigning exactly one expert per token (top-1 routing) across hundreds of TPU devices. Before this, MoE research existed but operated at much smaller scales, and the communication primitives required for expert parallelism were treated as exotic operations. The Switch Transformer paper brought all-to-all communication into the mainstream of large-model training and established the design patterns that later systems like GLaM, ST-MoE, and Mixtral would refine. The key insight was that the all-to-all communication cost, while real, was manageable compared to the compute savings from sparse activation, especially as models scaled to many billions of parameters.
The Distribution Challenge
To understand why expert parallelism is necessary, consider a concrete scenario. You are training an MoE model with 64 experts, where each expert is a two-layer feed-forward network with hidden dimension 4096. This is a common configuration: each expert acts as an independent specialist, capable of processing any token but trained to excel on a particular subset of the input space. With a 4x expansion factor in the feed-forward hidden dimension (a standard choice in transformer architectures), each expert consists of two linear transformations, one projecting up to 16384 dimensions and another projecting back down to 4096.
The parameter count scales rapidly with these choices. Each of the two linear layers in an expert contains the product of its input and output dimensions in weights. Summing across both layers, each expert holds approximately 134 million parameters. That number may sound manageable in isolation, but when you multiply by 64 experts, the expert layers alone require storing 8.6 billion parameters. And this is just the expert component: the model also has attention layers and layer normalization, plus embeddings and positional encodings, all of which add to the total. A realistic large-scale MoE model might have between 50 billion and several trillion total parameters, making the hardware requirements extraordinary.
Only 2 of the 64 experts activate for any given token in a top-2 routing scheme, but all 64 experts must exist in memory during training. You cannot materialize only the experts needed for a batch because gradient computation requires the entire parameter set to be accessible. The full 8.6 billion expert parameters must reside somewhere in the cluster's memory at all times during training. No single GPU can efficiently hold all of this. Even the NVIDIA H100 with 80 GB of HBM3 cannot hold it alone once optimizer states and activations are included alongside non-expert components. Distributing computation lets each device focus on a subset of experts and run expert computation in parallel, reducing per-step training time.
The mathematics of this situation follows directly from the architecture choices. Each expert's parameter count comes from two linear transformations: one that projects from the model dimension to the feed-forward hidden dimension, and another that projects back. When you multiply model dimension by the expanded hidden dimension for each layer and then by the number of experts, the total parameter count grows rapidly. For sufficiently large MoE models, expert parallelism is a hardware necessity rather than an optional optimization.
# Example MoE configuration
model_dim = 4096
ffn_hidden = 4096 * 4 # Standard 4x expansion
num_experts = 64
num_gpus = 8
# Parameters per expert (two linear layers)
params_per_expert = model_dim * ffn_hidden + ffn_hidden * model_dim
total_expert_params = params_per_expert * num_expertsParameters per expert: 134.2M Total expert parameters: 8.59B Experts per GPU (with expert parallelism): 8
With 8 GPUs, expert parallelism places 8 experts on each device. This is the core principle of the strategy: partition the expert pool across the available devices rather than replicating it. Instead of asking each device to maintain a copy of every expert, which would be both memory-prohibitive and wasteful of network bandwidth for gradient synchronization, we ask each device to host a specific subset of experts. Tokens that need those experts must travel to that device, but the memory burden is distributed evenly across the cluster. The router (gating network) is replicated on all devices because it is small and must be accessible everywhere to make routing decisions, but the experts themselves each live on exactly one device.
This design choice means that communication in expert parallelism differs fundamentally from other forms of parallelism. In data parallelism, communication happens during the backward pass when gradients are aggregated across replicas; the forward pass requires no inter-device communication. In expert parallelism, communication happens during the forward pass whenever a token needs to visit an expert on a different device. This makes expert parallelism sensitive to network bandwidth and prevents communication from being hidden behind computation as easily as gradient synchronization.
Expert Placement Strategies
Once we commit to distributing experts across devices, we face a non-trivial design decision: which experts should live on which devices? This seemingly simple question has meaningful implications for system performance, communication patterns, and load balance. The placement decision is made once at model initialization (or occasionally updated during training) and affects every single forward and backward pass. A good placement strategy can substantially reduce load imbalances that would otherwise cause some devices to idle while others are overwhelmed.
The simplest and most common strategy is uniform placement, where experts are evenly divided across devices. With experts and devices, each device holds experts. This approach requires no prior knowledge about how experts will be used and creates predictable, symmetric communication patterns. Every device participates in the all-to-all with the same amount of data sent and received (assuming balanced routing), which simplifies buffer allocation and makes performance analysis tractable. Uniform placement is the default choice in virtually all production MoE implementations, from Switch Transformer to Mixtral.
The appeal of uniform placement goes beyond simplicity. When expert utilization is well-balanced (which good load-balancing losses aim to ensure), uniform placement achieves optimal memory efficiency: each device holds exactly the same amount of expert parameters, and no device is a memory bottleneck. The predictability of the communication pattern also allows network schedulers to optimize traffic flows, which is not possible with irregular placement strategies. For these reasons, most MoE research and production systems default to uniform placement and invest effort in making routing balanced rather than trying to match placement to routing patterns.
def uniform_expert_placement(num_experts: int, num_devices: int) -> dict:
"""
Assign experts to devices uniformly.
Returns mapping from expert_id to device_id.
"""
assert num_experts % num_devices == 0, "Experts must divide evenly"
experts_per_device = num_experts // num_devices
placement = {}
for expert_id in range(num_experts):
device_id = expert_id // experts_per_device
placement[expert_id] = device_id
return placement
# Create placement for our example
placement = uniform_expert_placement(num_experts=64, num_devices=8)Expert placement (expert_id -> device_id): Device 0: experts 0-7 Device 1: experts 8-15 Device 2: experts 16-23 Device 3: experts 24-31 Device 4: experts 32-39 Device 5: experts 40-47 Device 6: experts 48-55 Device 7: experts 56-63

The placement mapping shown above illustrates the contiguous assignment strategy: experts 0 through 7 reside on device 0, experts 8 through 15 on device 1, and so forth. This contiguous assignment simplifies the mental model and makes it easy to compute which device holds any given expert through simple integer division: device_id = expert_id // experts_per_device. There is no lookup table required; the mapping is a pure arithmetic operation. This simplicity extends to the backward pass as well, since gradient updates for a given expert always flow to the same device, avoiding the need for cross-device gradient synchronization for expert parameters.
Uniform placement works well when load balancing is effective and all experts receive roughly equal traffic. However, as we discussed in the chapters on load balancing and auxiliary losses, achieving perfect balance is challenging. Some experts inevitably become more popular than others, partly due to the expressiveness of the gating network and partly due to natural variation in the input distribution. When experts are unevenly utilized, the device holding a popular expert becomes a computational bottleneck, while devices holding underutilized experts sit partially idle. This is the basic tension that uniform placement does not resolve on its own.
Capacity-Aware Placement
An alternative strategy accounts for expected expert utilization when making placement decisions. If certain experts consistently receive more tokens, placing them on separate devices can better balance computation. The intuition is straightforward: if experts 0 and 1 are both heavily used, placing them on the same device would overload it while leaving others underused. By spreading popular experts across different devices, we achieve better computational balance without changing anything about the routing algorithm itself.
Think of this strategy like a supermarket manager who knows that the bakery counter and the deli counter are both always busy. Rather than placing them next to each other (which would create a congested aisle), the manager positions them at opposite ends of the store, distributing foot traffic more evenly. The placement decision is made using knowledge of customer behavior, not customer identity; the manager does not need to know who specifically will visit the bakery, just that it will be popular.
Capacity-aware placement requires knowledge of expert utilization patterns, which are not available at initialization. In practice, you can estimate expected utilization from a pilot run or from the early stages of training, then make placement decisions based on those estimates. A common heuristic is to measure expert selection frequencies on a sample of training data before committing to a placement strategy. Alternatively, if you are continuing a training run from a checkpoint, you can use the utilization statistics from the previous run to inform placement in the next run.
The greedy bin-packing algorithm processes experts in descending order of expected load and assigns each one to the device with the lowest current load total. This is the classic First Fit Decreasing (FFD) heuristic, which is known to produce near-optimal solutions for bin-packing problems. While not guaranteed optimal, it reliably achieves load variance significantly lower than random or contiguous assignment when expert utilization is non-uniform.
from typing import Tuple
import torch
def capacity_aware_placement(
num_experts: int,
num_devices: int,
expert_loads: torch.Tensor, # Expected load per expert
) -> Tuple[dict, torch.Tensor]:
"""
Place experts to balance computational load across devices.
Greedy algorithm: assign each expert to the least-loaded device.
"""
device_loads = torch.zeros(num_devices)
placement = {}
# Sort experts by load (descending) for better packing
sorted_experts = torch.argsort(expert_loads, descending=True)
for expert_id in sorted_experts.tolist():
# Assign to device with minimum current load
target_device = torch.argmin(device_loads).item()
placement[expert_id] = target_device
device_loads[target_device] += expert_loads[expert_id]
return placement, device_loads
# Simulate non-uniform expert loads
torch.manual_seed(42)
expert_loads = torch.softmax(torch.randn(64), dim=0) * 64 # Normalized loads
placement_aware, final_loads = capacity_aware_placement(64, 8, expert_loads)Capacity-aware placement results: Load variance across devices: 0.0004 Uniform placement load variance: 5.0085


The greedy algorithm presented here processes experts in order of their expected load, always assigning each expert to the device with the lowest current total load. This bin-packing approach, while not optimal in all cases, achieves a meaningful reduction in load variance compared to uniform placement. The variance reduction shown in the output demonstrates that even this simple heuristic can significantly improve balance when expert utilization patterns are non-uniform. In practice, the improvement depends on how skewed the utilization distribution is: for nearly uniform distributions, the two strategies converge; for highly skewed distributions, capacity-aware placement can reduce per-device peak load substantially.
Capacity-aware placement reduces load imbalance but requires knowledge of expert utilization patterns, which may change during training. This creates a practical challenge: the routing patterns that determine expert popularity evolve as the model learns. An expert that starts as generalist may develop specialization over time, attracting more traffic from certain types of input. A placement decision made at initialization may therefore become suboptimal as training progresses. In principle, you could re-evaluate placement periodically during training, but this requires checkpointing, re-assigning experts (which may require moving parameters between devices), and restarting training, adding significant operational complexity. In practice, most implementations use uniform placement combined with strong load balancing losses, preferring to shape the routing behavior rather than adapt the physical placement.
All-to-All Communication
Expert parallelism relies on a communication primitive called all-to-all. This primitive is fundamentally different from the collective operations commonly used in dense model training, and understanding its behavior is needed for reasoning about MoE system performance. Getting a firm grip on all-to-all communication is arguably the most important conceptual step in understanding how expert parallelism works at the systems level, because it is where the theoretical elegance of MoE routing meets the hard constraints of physical hardware.
In dense model training, the dominant communication pattern is all-reduce: every device contributes a value (typically a gradient tensor), and all devices receive the sum of all contributions. This pattern is highly efficient because modern collective communication libraries like NCCL implement ring-based all-reduce algorithms that use every link in the network optimally. The bandwidth utilization approaches the theoretical maximum, and the latency scales only logarithmically with device count. All-reduce is a symmetric operation: every device sends the same amount of data and receives the same amount in return.
All-to-all is fundamentally asymmetric. Rather than every device contributing the same value to a shared sum, all-to-all involves every device sending a potentially different message to every other device. Each device has distinct pieces of data, one destined for each of the other devices (including itself), and the operation orchestrates all data transfers simultaneously. The total data volume is the same as in the all-reduce case, but the communication pattern is much harder to optimize because the data to be transferred varies between device pairs and depends on the dynamic routing decisions made in the forward pass.
Unlike all-reduce (which aggregates values across devices) or all-gather (which collects values from all devices), all-to-all performs a complete exchange where each device sends different data to each other device. Think of it as a postal system where every post office simultaneously sends letters to every other post office, and all the letters arrive at their destinations at roughly the same time. Each device has a different message for each recipient, and all messages are exchanged in a single coordinated operation.
An all-to-all operation takes input tensors partitioned across devices and redistributes them so that each device receives the portions destined for it from all other devices. If device has data for device , that data moves from to , and this happens simultaneously for all device pairs. The total volume of data in the system remains constant; only its distribution changes.
All-to-all communication fits into the MoE forward pass in four phases:
-
Routing decision: Each device computes which experts each of its tokens should visit. The gating network produces probabilities over all experts, and the top-k selection determines which experts will process each token.
-
First all-to-all (dispatch): Tokens are sent to the devices holding their target experts. A token that was generated on device 0 but needs expert 42 (located on device 5) must travel across the network to device 5.
-
Expert computation: Each device processes tokens using its local experts. At this point, each device has received all tokens that need its experts, regardless of where those tokens originated.
-
Second all-to-all (combine): Results are sent back to the original devices. The processed token representations must return to the devices where they originated so that the model can continue processing them through subsequent layers.
The two all-to-all operations are needed bookends around the expert computation. The first operation disperses tokens out to their target experts; the second gathers the results back. Without both operations, tokens processed on remote devices would never return to their home devices, and the model's residual stream would be incomplete. This two-operation structure means that every MoE layer incurs the communication cost twice, which is an important factor when calculating total communication overhead.
def simulate_all_to_all_routing(
tokens_per_device: int, num_devices: int, num_experts: int, top_k: int = 2
) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Simulate the routing decisions in expert parallelism.
Returns dispatch counts (tokens sent from device i to device j).
"""
experts_per_device = num_experts // num_devices
# Simulate random expert assignments for each token on each device
total_tokens = tokens_per_device * num_devices
# Each token selects top_k experts
expert_indices = torch.randint(
0, num_experts, (num_devices, tokens_per_device, top_k)
)
# Count how many tokens each device sends to each other device
dispatch_counts = torch.zeros(num_devices, num_devices, dtype=torch.long)
for src_device in range(num_devices):
for token_idx in range(tokens_per_device):
for k in range(top_k):
expert_id = expert_indices[src_device, token_idx, k].item()
dst_device = expert_id // experts_per_device
dispatch_counts[src_device, dst_device] += 1
return dispatch_counts, expert_indices
# Simulate routing for a batch
sim_tokens = 1024
sim_devices = 8
sim_experts = 64
sim_top_k = 2
dispatch_counts, _ = simulate_all_to_all_routing(
tokens_per_device=sim_tokens,
num_devices=sim_devices,
num_experts=sim_experts,
top_k=sim_top_k,
)Dispatch matrix (rows=source device, cols=destination device): Each entry shows token-expert pairs sent from source to destination [[250 224 248 272 286 230 270 268] [257 262 267 252 256 245 250 259] [253 234 275 249 264 262 261 250] [225 264 263 257 235 257 256 291] [246 256 284 252 228 264 268 250] [256 258 257 258 269 237 259 254] [243 293 264 252 258 257 228 253] [262 257 254 266 230 247 259 273]] Total token-expert pairs: 16384 Expected per device: 2048 (tokens × top_k)

The dispatch matrix reveals the communication pattern in concrete terms. Each row represents a source device, and each column represents a destination device. The value at position (i, j) tells us how many token-expert pairs device i sends to device j during the dispatch phase. With uniform random routing and 8 devices, each device sends roughly equal traffic to all devices (including itself). The values along the diagonal are particularly interesting: these represent tokens that happen to route to experts on their original device, requiring no network communication. These "local" routings are essentially free from a communication perspective, though they still require computation.
The All-to-All Operation
PyTorch's distributed library provides all_to_all for this communication pattern. The operation coordinates all devices to simultaneously exchange data according to their individual send and receive specifications. This coordination is complex at scale, as it requires synchronization across all participating devices. In practice, the operation is implemented using ring-based algorithms or switch-based collective communication that can exploit the topology of the underlying network hardware.
A important detail is that the all-to-all operation in expert parallelism involves variable-length messages: device i might send 1000 token representations to device j but only 200 to device k, depending on how the gating network has routed tokens for this particular batch. This variability makes the all-to-all in expert parallelism more complex than the uniform all-to-all that arises in some other distributed algorithms. Most production frameworks handle this through a variant called all_to_all_single (in PyTorch) or all_to_allv (in MPI), which accepts explicit send and receive counts for each device pair.
def expert_parallel_dispatch(
hidden_states: torch.Tensor, # [batch_size, seq_len, hidden_dim]
expert_indices: torch.Tensor, # [batch_size, seq_len, top_k]
expert_weights: torch.Tensor, # [batch_size, seq_len, top_k]
num_experts: int,
world_size: int,
rank: int,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""
Dispatch tokens to their target devices via all-to-all.
This is a simplified version showing the key operations.
Production implementations handle variable-length messages.
"""
batch_size, seq_len, hidden_dim = hidden_states.shape
experts_per_device = num_experts // world_size
# Flatten to [num_tokens, hidden_dim]
tokens = hidden_states.view(-1, hidden_dim)
num_tokens = tokens.shape[0]
top_k = expert_indices.shape[-1]
# Determine destination device for each expert selection
flat_indices = expert_indices.view(-1) # [num_tokens * top_k]
dest_devices = flat_indices // experts_per_device
# Sort tokens by destination device for efficient all-to-all
sort_indices = torch.argsort(dest_devices)
# Count tokens going to each device
send_counts = torch.bincount(dest_devices, minlength=world_size)
# In actual implementation, we'd call:
# recv_tokens = torch.empty(...)
# dist.all_to_all_single(recv_tokens, send_tokens, recv_counts, send_counts)
return tokens, send_counts, sort_indices
# Demonstrate dispatch logic
demo_states = torch.randn(1, 4, 4)
# Experts mapped to devices: 0-1->Dev0, 2-3->Dev1
demo_experts = torch.tensor([[[0, 2], [1, 3], [0, 1], [2, 3]]])
demo_weights = torch.ones_like(demo_experts, dtype=torch.float)
_, demo_counts, _ = expert_parallel_dispatch(
hidden_states=demo_states,
expert_indices=demo_experts,
expert_weights=demo_weights,
num_experts=4,
world_size=2,
rank=0,
)Dispatch counts (2 devices): [4, 4]
The dispatch counts demonstrate the core accounting that drives the all-to-all operation. In this small example with 4 tokens and 2 experts per token, we see that 4 token-expert pairs are destined for device 0 and 4 for device 1. This reflects a balanced routing pattern across the two devices. In a real distributed setting, these counts would be exchanged between devices before the all-to-all. This allows each device to allocate the correct amount of receive buffer space. This pre-exchange of counts adds a small but necessary synchronization step to the dispatch phase.
The key insight here is that all-to-all enables dynamic routing: each token goes exactly where it needs to go based on the gating network's decision. This differs fundamentally from static parallelism strategies where data movement patterns are fixed at compile time. In tensor parallelism, for example, the communication pattern is determined entirely by the model architecture and is identical for every input. In expert parallelism, the communication pattern depends on the content of the input itself, as mediated by the learned gating network. This dynamic behavior is both the source of MoE's flexibility and the root of its systems complexity. You cannot pre-allocate communication buffers based on a fixed pattern; you must compute the routing, determine send counts, communicate those counts, allocate buffers, and then perform the actual data transfer.
Communication Overhead Analysis
Expert parallelism introduces communication costs that do not exist in dense models. Understanding these costs is needed for system design and performance optimization. While dense model training also involves communication (for gradient synchronization, for example), the all-to-all operations in MoE occur in the forward pass itself, directly adding latency to every training step. Unlike gradient synchronization, which can often be overlapped with the backward pass computation through techniques like gradient bucketing, the dispatch and combine all-to-all operations in the forward pass sit on the necessary path: the expert computation cannot begin until the dispatch is complete, and the next layer cannot begin until the combine is complete.
This is a subtle but important distinction. In data-parallel training with gradient synchronization, the communication is largely "free" from a latency perspective because it happens in the background while computation continues. In expert parallelism, the forward-pass communication creates unavoidable serial dependencies. The training step time is therefore: routing computation, plus dispatch all-to-all, plus expert computation, plus combine all-to-all, plus remaining layer computation. The all-to-all operations add directly to the necessary path length, not just to bandwidth utilization.
Communication Volume
For each MoE layer, we perform two all-to-all operations. The first operation dispatches tokens from their originating devices to the devices holding their target experts. The second operation returns the processed results back to the originating devices. Each operation moves the same amount of data, as every token that goes out must come back. The total communication volume is therefore symmetric: dispatch and combine are mirror images of each other.
The total communication volume (in elements) per MoE layer can be expressed in terms of the key configuration parameters. We compute this by multiplying the number of token-expert pairs by the size of each token representation, then doubling for the two communication phases:
where:
- : the batch size (number of sequences). Larger batches increase throughput but also increase the total number of tokens that must be routed.
- : the sequence length (number of tokens per sequence). Longer sequences mean more tokens per batch, directly scaling communication.
- : the top- parameter (number of experts per token). When , each token creates two communication events rather than one.
- : the hidden dimension (size of each token vector). Larger models with wider hidden dimensions must send more data per token.
- : accounts for the two all-to-all operations (dispatch and combine). This factor is fixed by the algorithm structure.
Why does this formula make sense? Notice that gives the total number of tokens in the batch. Multiplying by gives the total number of token-expert pairs, since each token creates routing events. Multiplying by converts token count to element count, since each token is a vector of elements. The factor of 2 reflects that each element travels twice: once to the expert device and once back. The multiplicative relationship between these factors means that communication volume grows rapidly with any of the individual parameters. Doubling any single factor doubles the communication cost.
def calculate_communication_volume(
batch_size: int,
seq_len: int,
hidden_dim: int,
top_k: int,
num_moe_layers: int,
bytes_per_element: int = 2, # bf16
) -> dict:
"""
Calculate total communication volume for MoE forward pass.
"""
# Tokens communicated per MoE layer
tokens_per_layer = batch_size * seq_len * top_k
elements_per_layer = tokens_per_layer * hidden_dim
# Two all-to-all operations per layer (dispatch + combine)
elements_per_layer_total = elements_per_layer * 2
# Total across all MoE layers
total_elements = elements_per_layer_total * num_moe_layers
total_bytes = total_elements * bytes_per_element
return {
"tokens_per_layer": tokens_per_layer,
"elements_per_all_to_all": elements_per_layer,
"total_elements": total_elements,
"total_bytes": total_bytes,
"total_gb": total_bytes / (1024**3),
}
# Typical training configuration
comm_volume = calculate_communication_volume(
batch_size=32,
seq_len=2048,
hidden_dim=4096,
top_k=2,
num_moe_layers=16, # MoE in every other layer of 32-layer model
)Communication volume analysis: Tokens per MoE layer: 131,072 Elements per all-to-all: 536,870,912 Total elements (forward pass): 17,179,869,184 Total communication: 32.00 GB
This substantial volume of data movement, roughly 32 GB per forward pass for a single batch, highlights why network bandwidth is often the bottleneck in MoE training. To put this in perspective, even with high-bandwidth interconnects capable of 200 GB/s, transferring 32 GB takes around 160 milliseconds if the transfers were purely sequential. In practice, the all-to-all implementation uses bidirectional transfers and pipeline overlap to reduce effective latency, but the basic constraint remains: you cannot transfer 32 GB faster than the interconnect bandwidth allows. The communication cost scales linearly with the number of MoE layers, making efficient routing needed. This is one reason why some architectures use MoE layers only in alternating positions rather than in every layer, reducing the number of all-to-all operations by half at the cost of slightly reduced expert capacity.
Communication vs Computation Trade-off
The key metric for understanding expert-parallelism efficiency is the ratio of communication time to computation time. If communication dominates, adding more devices provides diminishing returns because each device spends more time waiting for data than computing. Conversely, if computation dominates, the system can efficiently use additional devices for expert computation without communication becoming a bottleneck.
This ratio is determined by what is sometimes called the arithmetic intensity of the all-to-all: the ratio of compute operations performed to bytes transferred. For expert parallelism, we transfer token representations to expert devices and get back transformed representations. The computation performed on those representations is the expert forward pass, which involves two matrix multiplications per token. The ratio of FLOP to bytes transferred determines whether the system is compute-bound or bandwidth-bound.
As we scale to more devices, the fraction of tokens that must cross device boundaries increases. With devices, approximately of all token-expert pairs must be transferred across the network (assuming uniform routing). As grows, this fraction approaches 1, meaning almost all communication is cross-device. This is why communication overhead grows with device count: more devices means more cross-device traffic for the same total token count.
def analyze_scaling_efficiency(
hidden_dim: int,
ffn_hidden: int,
tokens_per_device: int,
top_k: int,
num_devices_list: list,
bandwidth_gb_per_sec: float = 200, # NVLink bandwidth
tflops: float = 150, # GPU compute throughput
) -> dict:
"""
Analyze how communication overhead scales with device count.
"""
results = []
for num_devices in num_devices_list:
# Communication: tokens × hidden_dim × 2 (dispatch + combine)
# Each token sends/receives from other devices
# Fraction of traffic that crosses device boundaries
cross_device_fraction = (num_devices - 1) / num_devices
comm_elements = (
tokens_per_device * top_k * hidden_dim * 2 * cross_device_fraction
)
comm_bytes = comm_elements * 2 # bf16
comm_time = comm_bytes / (bandwidth_gb_per_sec * 1e9)
# Computation: expert forward pass
# FLOPs = 2 * tokens * hidden_dim * ffn_hidden * 2 (two linear layers)
flops = 2 * tokens_per_device * top_k * hidden_dim * ffn_hidden * 2
compute_time = flops / (tflops * 1e12)
efficiency = compute_time / (compute_time + comm_time)
results.append(
{
"num_devices": num_devices,
"comm_time_ms": comm_time * 1000,
"compute_time_ms": compute_time * 1000,
"efficiency": efficiency,
}
)
return results
scaling_results = analyze_scaling_efficiency(
hidden_dim=4096,
ffn_hidden=4096 * 4,
tokens_per_device=4096,
top_k=2,
num_devices_list=[2, 4, 8, 16, 32, 64],
bandwidth_gb_per_sec=25, # InfiniBand cross-node bandwidth
)Expert parallelism scaling efficiency:
Devices Comm (ms) Compute (ms) Efficiency
--------------------------------------------------
2 2.684 14.660 84.5%
4 4.027 14.660 78.5%
8 4.698 14.660 75.7%
16 5.033 14.660 74.4%
32 5.201 14.660 73.8%
64 5.285 14.660 73.5%

Several important insights emerge from this analysis. First, communication overhead grows with the number of devices because more tokens must cross device boundaries. With 2 devices, half the tokens (on average) route to local experts and require no network transfer. With 64 devices, only 1/64 of tokens stay local, meaning the large majority must travel across the network. The cross-device fraction converges to 1 as device count increases, meaning the communication overhead asymptotes at a value determined solely by the bandwidth and the per-token message size.
Second, efficiency remains high when computation dominates, but degrades as we scale to many devices. The efficiency percentages in the table show this degradation clearly: at small device counts the system is nearly compute-bound, but at 32 or 64 devices the communication overhead becomes a significant fraction of total step time. This suggests a practical limit to how far expert parallelism alone can scale before communication becomes the bottleneck. In practice, reaching 64-device expert parallelism with good efficiency requires either extremely fast interconnects (like NVLink at 600+ GB/s) or expert dimensions large enough that computation dominates.
Third, the choice of interconnect has an outsized impact on achievable scale. The analysis above assumes 25 GB/s cross-node InfiniBand bandwidth, which is representative of typical GPU clusters. Higher-bandwidth interconnects like NVLink (used for intra-node GPU communication at up to 600 GB/s on modern systems) can push the efficiency curves significantly upward, letting larger device counts before communication becomes a bottleneck. This is why researchers at organizations with access to tightly-coupled GPU clusters (connected by NVSwitch or similar high-bandwidth fabrics) can afford to use larger expert parallelism configurations than teams relying on commodity networking.
Expert Parallelism Implementation
A complete expert parallelism implementation requires coordinating routing decisions, all-to-all communication, and expert computation. The implementation must handle several subtle challenges: variable-length messages between device pairs, preservation of token ordering for correct gradient computation, efficient batching of expert computations, and proper integration with the rest of the transformer architecture. Here is a simplified but functional implementation that demonstrates the core concepts without the complexity of a production distributed training framework:
from typing import Tuple
import torch
import torch.nn as nn
class ExpertParallelMoE(nn.Module):
"""
MoE layer with expert parallelism.
Each device holds a subset of experts.
"""
def __init__(
self,
hidden_dim: int,
ffn_hidden: int,
num_experts: int,
num_local_experts: int,
top_k: int = 2,
capacity_factor: float = 1.25,
):
super().__init__()
self.hidden_dim = hidden_dim
self.num_experts = num_experts
self.num_local_experts = num_local_experts
self.top_k = top_k
self.capacity_factor = capacity_factor
# Router (replicated on all devices)
self.gate = nn.Linear(hidden_dim, num_experts, bias=False)
# Local experts only
self.experts = nn.ModuleList(
[
nn.Sequential(
nn.Linear(hidden_dim, ffn_hidden, bias=False),
nn.GELU(),
nn.Linear(ffn_hidden, hidden_dim, bias=False),
)
for _ in range(num_local_experts)
]
)
def forward(
self, hidden_states: torch.Tensor, expert_parallel_group=None
) -> Tuple[torch.Tensor, dict]:
"""
Forward pass with expert parallelism.
Args:
hidden_states: [batch, seq_len, hidden_dim]
expert_parallel_group: Distributed process group for all-to-all
Returns:
output: [batch, seq_len, hidden_dim]
aux_info: Dictionary with auxiliary information
"""
batch_size, seq_len, hidden_dim = hidden_states.shape
num_tokens = batch_size * seq_len
# Compute routing probabilities
router_logits = self.gate(
hidden_states
) # [batch, seq_len, num_experts]
router_probs = torch.softmax(router_logits, dim=-1)
# Select top-k experts
top_k_probs, top_k_indices = torch.topk(
router_probs, self.top_k, dim=-1
) # [batch, seq_len, top_k]
# Normalize weights
top_k_weights = top_k_probs / top_k_probs.sum(dim=-1, keepdim=True)
# Flatten for processing
flat_hidden = hidden_states.view(num_tokens, hidden_dim)
flat_indices = top_k_indices.view(num_tokens, self.top_k)
flat_weights = top_k_weights.view(num_tokens, self.top_k)
# Process through experts (simplified: no actual distributed communication)
# In production, this would involve all-to-all operations
output = self._local_expert_forward(
flat_hidden, flat_indices, flat_weights
)
return output.view(batch_size, seq_len, hidden_dim), {
"router_probs": router_probs,
"top_k_indices": top_k_indices,
}
def _local_expert_forward(
self,
hidden_states: torch.Tensor,
expert_indices: torch.Tensor,
expert_weights: torch.Tensor,
) -> torch.Tensor:
"""
Process tokens through local experts.
Simplified version without capacity limits.
"""
num_tokens = hidden_states.shape[0]
output = torch.zeros_like(hidden_states)
# Determine which experts are local
world_size = self.num_experts // self.num_local_experts
rank = 0 # Would be dist.get_rank() in distributed setting
local_expert_start = rank * self.num_local_experts
local_expert_end = local_expert_start + self.num_local_experts
for k in range(self.top_k):
for local_idx, expert in enumerate(self.experts):
global_idx = local_expert_start + local_idx
# Find tokens routed to this expert
mask = expert_indices[:, k] == global_idx
if mask.sum() == 0:
continue
# Process tokens
expert_input = hidden_states[mask]
expert_output = expert(expert_input)
# Weight and accumulate
weights = expert_weights[mask, k : k + 1]
output[mask] += weights * expert_output
return output
# Create model instance
moe_layer = ExpertParallelMoE(
hidden_dim=512,
ffn_hidden=2048,
num_experts=8,
num_local_experts=8, # All experts local for this demo
top_k=2,
)# Test forward pass
test_input = torch.randn(2, 16, 512) # batch=2, seq_len=16
output, aux_info = moe_layer(test_input)Input shape: torch.Size([2, 16, 512]) Output shape: torch.Size([2, 16, 512]) Router probabilities shape: torch.Size([2, 16, 8]) Selected experts shape: torch.Size([2, 16, 2])
The output shapes confirm that the model processes the input sequence and produces an output of the same dimensionality. The router probabilities tensor has shape corresponding to the batch size, sequence length, and number of experts. This provides a complete probability distribution over experts for each token. The selected indices tensor has the same batch and sequence dimensions but with the last dimension equal to k, showing which experts were chosen for each token. These auxiliary outputs provide visibility into the gating mechanism. This allows us to verify routing behavior, compute load balancing metrics, and debug potential issues.
Importantly, notice that the gate (gating network) is a nn.Linear layer replicated across all devices, while the self.experts module list contains only the local subset of experts. This asymmetry is basic to the design: the routing decision must be made globally (every device must know which expert each of its tokens goes to), so the gate must be available everywhere. But the expert parameters are partitioned, with each device holding only its assigned subset. In a fully distributed implementation, gradients for the gate would be synchronized using all-reduce across all devices (since it is replicated), while gradients for local expert parameters would not require any inter-device communication (since each expert is owned by exactly one device).


Handling Variable-Length Communication
A subtlety of expert parallelism is that all-to-all messages have variable lengths. One device might send 1000 tokens to device 0 but only 500 to device 1, depending on routing decisions. This variability arises naturally from the learned gating function: different inputs activate different experts, and the distribution of selected experts changes from batch to batch. Even with perfect load-balancing losses, the exact counts per device pair will vary due to the discrete nature of top-k selection. This requires careful buffer management to ensure that receiving devices allocate sufficient space without wasting memory.
The standard approach is to exchange send counts before the all-to-all data transfer. Each device broadcasts its intended send sizes to all other devices, letting each receiver to compute exactly how much space it needs. This pre-exchange is itself a collective communication operation (an all-to-all of integers rather than token representations), but it is cheap because it involves only integers rather than floating-point values. Once every device knows how many tokens it will receive from each sender, it allocates the receive buffer and initiates the all-to-all data transfer.
The permutation indices used to restore token ordering after the combine step are equally important. After receiving processed results back from expert devices, the tokens arrive in an order determined by the second all-to-all's communication schedule, not necessarily by the original sequence positions. Without restoring the original order, the subsequent transformer layers would process tokens in the wrong positions, breaking the correspondence between token representations and their positions in the residual stream. The inverse permutation computed during the dispatch phase is a "return address" that guides the reordering step.
def prepare_all_to_all_buffers(
hidden_states: torch.Tensor,
expert_indices: torch.Tensor,
num_devices: int,
experts_per_device: int,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""
Prepare send/receive buffers for variable-length all-to-all.
Returns:
send_buffer: Sorted tokens ready for dispatch
send_counts: Number of elements to send to each device
permutation: Indices to restore original order after combine
"""
num_tokens, hidden_dim = hidden_states.shape
top_k = expert_indices.shape[1]
# Expand hidden states for each expert selection
expanded_hidden = hidden_states.unsqueeze(1).expand(-1, top_k, -1)
expanded_hidden = expanded_hidden.reshape(-1, hidden_dim)
# Determine destination device for each token-expert pair
flat_indices = expert_indices.view(-1)
dest_devices = flat_indices // experts_per_device
# Sort by destination device
sort_indices = torch.argsort(dest_devices, stable=True)
sorted_hidden = expanded_hidden[sort_indices]
# Count elements for each destination
send_counts = torch.bincount(dest_devices, minlength=num_devices)
# Create inverse permutation for restoration
inverse_indices = torch.empty_like(sort_indices)
inverse_indices[sort_indices] = torch.arange(len(sort_indices))
return sorted_hidden, send_counts, inverse_indices
# Test buffer preparation
num_test_tokens = 32
test_top_k = 2
test_hidden = torch.randn(num_test_tokens, 512)
test_experts = torch.randint(0, 64, (num_test_tokens, test_top_k))
send_buffer, send_counts, permutation = prepare_all_to_all_buffers(
test_hidden, test_experts, num_devices=8, experts_per_device=8
)Original hidden states: torch.Size([32, 512]) Send buffer (sorted): torch.Size([64, 512]) Send counts per device: [8, 3, 8, 8, 8, 9, 10, 10] Total dispatched: 64 (expected: 64) Permutation indices shape: torch.Size([64])
The buffer preparation function performs several important transformations. First, it expands the hidden states to account for the top-k selection, creating a separate copy of each token for each of its selected experts. Then, it sorts these expanded tokens by their destination device, grouping all tokens bound for device 0 together, followed by all tokens bound for device 1, and so on. This sorting is needed because the all-to-all operation expects contiguous chunks of data for each destination, and mixing tokens from different destinations in memory would require additional scatter-gather operations.
The permutation indices are important: after receiving results back via the second all-to-all, we use them to restore the original token ordering. Without this restoration step, the processed representations would be scrambled, breaking the correspondence between tokens and their positions in the sequence. The inverse permutation ensures that after the round-trip through expert computation, each token's result ends up in the correct position for subsequent layers. In practice, this permutation-inverse permutation pattern appears repeatedly in high-performance MoE implementations, and it is often one of the first places to look when debugging correctness issues in distributed MoE training.
Worked Example: Tracing a Token Through Expert Parallelism
To make the concepts concrete, let us trace a single token through the full expert-parallel MoE layer step by step. Suppose we have 4 devices, each holding 4 experts (16 experts total). Our token starts on device 0 and uses top-2 routing.
Step 1: Routing decision. The gating network on device 0 processes the token's hidden representation (a vector of dimension ) and produces logits for all 16 experts. After softmax normalization, the top-2 selection identifies expert 2 (local, on device 0) and expert 11 (remote, on device 2) as the chosen experts. The routing weights are, say, 0.6 and 0.4 respectively.
Step 2: Dispatch phase. The token creates two dispatch entries: one for expert 2 (destination: device 0) and one for expert 11 (destination: device 2). The dispatch entry for expert 2 is a "local" routing that does not require network transfer. The dispatch entry for expert 11 must travel across the network. During the all-to-all, device 0 sends a copy of this token's representation to device 2 along with a tag identifying the target expert (expert 11, which is local expert index 3 on device 2, computed as ). At the same time, device 0 may receive token representations from other devices that need to be processed by experts 0-3.
Step 3: Expert computation. Device 0 processes all tokens routed to experts 0-3, including the token with expert 2. Device 2 processes all tokens routed to experts 8-11, including the copy of our token for expert 11. Each device processes its tokens independently and in parallel, with no inter-device communication during this phase. The expert transforms the token representation through its two-layer feed-forward network.
Step 4: Combine phase. Device 2 sends the processed result for expert 11 back to device 0. During the second all-to-all, all processed results travel back to their originating devices. Device 0 now has two processed representations: the output from expert 2 (which never left device 0) and the output from expert 11 (returned from device 2). The final output is the weighted sum: .
Step 5: Residual connection. The combined output is added to the original token representation via the residual connection, and the result flows to the next transformer layer. From this layer's perspective, the computation looks identical to a dense feed-forward layer: input goes in, output comes out. The distributed routing machinery is completely transparent to subsequent layers.
This trace illustrates several key design choices. The routing weights (0.6 and 0.4) determine how much each expert contributes to the final representation, letting the model to blend expert knowledge. The local routing for expert 2 is free from a communication standpoint, which is why some implementations add incentives for local routing (though this can hurt model quality if taken too far). And the two-phase all-to-all structure ensures that every token gets both its expert results without the ordering information being lost.
Key Parameters
The key parameters for the Expert Parallel MoE implementation shape both performance and model quality. Understanding what each parameter controls helps you make informed choices when designing a new MoE configuration or debugging an existing one.
The first parameter to consider is num_experts, the total number of experts in the model. This determines the potential for specialization and affects both model capacity and communication patterns. Common values range from 8 (as in Mixtral-8x7B) to hundreds or even thousands (as in Switch Transformer and its descendants). More experts provide greater capacity and finer-grained specialization opportunities, but also increase the probability of routing imbalances and the complexity of the all-to-all communication pattern.
The parameter num_local_experts specifies how many experts reside on the current device, typically num_experts / world_size. This value determines the memory footprint per device and must divide evenly into the total expert count for uniform placement. If your cluster has 8 GPUs and you want 64 experts, each GPU holds 8 local experts. This ratio, the number of experts per device, also affects the expected fraction of local routing: with local experts out of total, on average of all token-expert pairs route locally without requiring network communication.
The top_k parameter, typically 1 or 2 in production models, controls how many experts each token visits. The Switch Transformer used top-1 to minimize communication overhead; subsequent work showed that top-2 provides better model quality at a modest increase in communication cost. Values above 2 are rare in practice because the marginal quality improvement diminishes while the communication cost grows linearly. Some recent architectures explore more complex routing schemes, such as always routing to one "shared" expert and using top-k for the remaining selections, which provides a floor of universal computation while retaining selective specialization.
The capacity_factor parameter reserves extra buffer space for load imbalances, typically set to values between 1.0 and 1.5. A value of 1.25 means each expert is pre-allocated 25% more capacity than the expected average load. This headroom absorbs natural variation in routing patterns, preventing token dropping when some experts receive more traffic than expected in a given batch. The tradeoff is memory: higher capacity factors consume more memory per expert, reducing the maximum batch size you can use. When tokens exceed expert capacity, they are typically dropped (their contribution to the output is zeroed), creating a training signal artifact that can harm model quality if the dropping rate is high.
Combining Expert Parallelism with Other Strategies
In practice, expert parallelism is combined with other parallelism strategies to scale MoE models to hundreds of GPUs. No single parallelism strategy can address all the challenges of training massive models: expert parallelism handles the distribution of experts, but other strategies are needed to handle large batch sizes, oversized individual layers, and the sequential dependencies between layers. Understanding how to compose these strategies is needed for anyone building or deploying large-scale MoE systems.
The three main parallelism strategies in deep learning are data parallelism, tensor parallelism, and pipeline parallelism. Each addresses a different bottleneck: data parallelism increases throughput by processing more data simultaneously, tensor parallelism reduces per-layer memory by sharding weight matrices, and pipeline parallelism enables sequential computation on very deep networks. Expert parallelism is a fourth orthogonal strategy that specifically addresses the memory and compute challenges of MoE architectures. In large-scale training, all four strategies are often combined simultaneously, creating what practitioners call "4D parallelism."
-
Expert + Data Parallelism: Replicate the expert parallel setup across multiple groups of devices. Each group has a complete set of experts, and different groups process different batches. This combination multiplies throughput by the number of data parallel replicas while keeping the communication complexity of expert parallelism contained within each replica group. The all-to-all communication happens within replica groups, while gradient synchronization happens across replica groups. This separation is beneficial because it allows the two communication patterns to use different network topologies: the all-to-all benefits from low-latency, high-bandwidth links within a node, while gradient synchronization can tolerate higher latency on cross-node connections.
-
Expert + Tensor Parallelism: Shard individual experts across devices in addition to distributing experts. This handles cases where individual experts are too large for a single device. For example, if each expert has 7 billion parameters (as in Mixtral-scale models), tensor parallelism can split each expert across 2 or more devices, with expert parallelism then distributing these sharded experts across additional device groups. The combination introduces additional all-reduce operations within each sharded expert, but allows training models with extremely large experts that would not fit on any single device.
-
Expert + Pipeline Parallelism: Distribute different layers across different pipeline stages, with expert parallelism within each MoE layer. This approach addresses the sequential nature of transformer computation, where each layer depends on the output of the previous layer. Pipeline parallelism overlaps computation across stages by processing multiple microbatches simultaneously, while expert parallelism handles the within-layer distribution of experts. The challenge is that the all-to-all operations within each MoE layer create synchronization points that can reduce pipeline efficiency, particularly if the pipeline is not carefully designed to overlap these operations with other computation.
import math
def calculate_device_requirements(
num_experts: int,
expert_params: int,
model_params: int, # Non-expert parameters
memory_per_device_gb: float,
bytes_per_param: int = 2, # bf16
) -> dict:
"""
Calculate minimum devices needed for different parallelism strategies.
"""
total_expert_memory_gb = (
num_experts * expert_params * bytes_per_param
) / 1e9
total_model_memory_gb = (model_params * bytes_per_param) / 1e9
total_memory_gb = total_expert_memory_gb + total_model_memory_gb
# Account for optimizer states and activations (rough 3x multiplier)
training_memory_gb = total_memory_gb * 3
# Expert parallelism only
min_ep_devices = max(
1, math.ceil(total_expert_memory_gb / memory_per_device_gb)
)
# With tensor parallelism for non-expert params
min_combined_devices = max(
math.ceil(total_expert_memory_gb / memory_per_device_gb),
math.ceil(total_model_memory_gb / memory_per_device_gb),
)
return {
"total_expert_memory_gb": total_expert_memory_gb,
"total_model_memory_gb": total_model_memory_gb,
"training_memory_gb": training_memory_gb,
"min_expert_parallel_devices": min_ep_devices,
"min_combined_devices": min_combined_devices,
}
# Mixtral-8x7B scale model
requirements = calculate_device_requirements(
num_experts=8,
expert_params=7_000_000_000, # 7B params per expert
model_params=7_000_000_000, # Shared params (attention, embeddings)
memory_per_device_gb=80, # A100 80GB
)Device requirements analysis (Mixtral-8x7B scale): Expert memory: 112.0 GB Other model memory: 14.0 GB Training memory (with optimizer): 378.0 GB Minimum devices (expert parallel): 2

The analysis shows that for a model of this scale, the expert parameters alone (112 GB) exceed the memory of a single 80GB device. We require at least 2 devices just to store the experts, and significantly more when training due to optimizer states and activations. The 3x multiplier for training memory accounts for the Adam optimizer (which maintains two additional tensors per parameter: the first and second moment estimates) and the activation checkpointing overhead. The total training memory requirement of over 350 GB means that even a fully populated 8-GPU server with 640 GB of total GPU memory may not be sufficient for training without careful memory management through techniques like gradient checkpointing, CPU offloading, or mixed precision training. In practice, teams training models at this scale often use 8, 16, or more devices, combining expert parallelism with data parallelism to achieve both memory capacity and training throughput.
Limitations and Impact
Expert parallelism enables scaling MoE models that would otherwise be impossible to train, but it comes with significant challenges that you must understand and address when designing real systems. These limitations are not theoretical: they represent concrete engineering problems that affect training stability, hardware efficiency, and model quality in practice.
The most immediate limitation is communication overhead. The all-to-all operations required for expert dispatch and combine add latency to every training step, directly on the necessary path. Models like Switch Transformer sometimes use expert parallelism across hundreds of devices, where communication costs become substantial. Techniques like capacity factors and token dropping help manage this by limiting the maximum number of tokens each expert must process, but they introduce a different problem: dropped tokens represent information the model never processes, creating a subtle but real degradation in the quality of the learned representations. A model that systematically drops 5% of tokens due to capacity constraints is effectively training on incomplete data, with unpredictable effects on downstream performance. The tradeoff between communication efficiency and token dropping is one of the central engineering challenges in MoE training.
Load imbalance creates another persistent challenge that interacts with both communication efficiency and model quality. Even with auxiliary losses encouraging balanced routing, some experts inevitably receive more traffic than others during training. This asymmetry means some devices finish their computation earlier and must wait for the slower devices to complete before the combine phase can begin. The problem worsens with more devices: a system with 64 devices has 64 opportunities for one device to become a bottleneck, compared to just 2 opportunities with 2 devices. The expected maximum load (the load on the busiest device) grows as approximately times the average load for devices with random routing, meaning load imbalance becomes a more severe problem as the system scales. Strong load-balancing losses can mitigate this, but they may constrain the routing function in ways that prevent effective expert specialization.
Communication topology sensitivity is a less obvious but equally important limitation. The efficiency of all-to-all communication depends critically on the underlying network topology. On a single node with NVLink, the all-to-all can proceed at very high bandwidth because all GPUs are directly connected. Across nodes connected by InfiniBand or Ethernet, the available bandwidth per device pair drops dramatically, and the communication overhead can dwarf the computation time. This means that expert parallelism configurations that work efficiently within a single node may become communication-bound when scaled across a multi-node cluster. Practical systems often restrict expert parallelism to within a single node and use data parallelism across nodes, precisely to avoid this topology mismatch.
Despite these challenges, expert parallelism has enabled larger language models. It allows models to have many more parameters than dense architectures while maintaining reasonable compute costs. The Switch Transformer demonstrated training models with trillions of parameters using expert parallelism, and follow-on research refined the approach. More recent models like Mixtral use expert parallelism more conservatively, with 8 experts balancing capacity and communication efficiency. The 8-expert configuration is popular because it maps cleanly to 8-GPU nodes, allowing expert communication over NVLink without cross-node traffic.
The impact extends beyond just scale. Expert parallelism enables a form of model specialization that dense models cannot achieve. Different experts can learn different capabilities, and the routing mechanism determines which capabilities to apply for each input. This architectural choice influences training efficiency and the kinds of representations these models learn. Empirical analyses of trained MoE models have found that experts tend to specialize in meaningful ways: some focus on particular syntactic patterns (function words versus content words), others develop domain-specific competence (code versus natural language versus scientific text), and some specialize by language in multilingual models. This emergent specialization is only possible because expert parallelism makes it practical to maintain many separate expert networks within a single model, giving each expert the capacity to develop a distinct representational profile without the constraints that would arise from sharing parameters.
The interaction between expert parallelism and model quality is an active area of research. Questions remain about the optimal number of experts for a given compute budget, how to choose the right capacity factor, whether top-1 or top-2 routing provides better quality-efficiency tradeoffs, and how to design load-balancing objectives that improve behavior rather than merely penalize imbalance. As hardware evolves and new communication technologies emerge, the sweet spot for expert-parallelism configurations will continue to shift, requiring updated empirical and theoretical analysis.
Summary
Expert parallelism is the distributed computing strategy that makes large-scale MoE models practical. Rather than replicating all experts on every device (which is memory-prohibitive) or sharding experts arbitrarily (which creates complex parameter ownership), expert parallelism assigns complete experts to devices and uses all-to-all communication to route tokens to wherever their selected experts reside. The strategy works because expert computation is independent across experts: once a device receives all tokens destined for its experts, it can process them locally without further coordination, and the results can be returned asynchronously via the second all-to-all.
The core concepts work together in a pipeline: expert placement assigns experts to devices at initialization, setting up the mapping that routing decisions will use. The gating network makes routing decisions in the forward pass, creating expert indices for each token. All-to-all dispatch sends tokens to their target devices; expert computation processes them locally; all-to-all combine returns results to originating devices. The entire pipeline repeats for every MoE layer in the model and every batch during training.
The key considerations for designing an expert-parallel system include:
- Expert placement assigns experts to devices, typically uniformly, though capacity-aware placement can improve load balance when expert utilization patterns are known or can be estimated from prior runs.
- All-to-all communication is the primitive that enables dynamic routing, where each device sends different data to each other device based on the gating network's decisions. The two all-to-all operations per MoE layer sit on the necessary path and cannot be hidden behind computation.
- Communication overhead grows with device count and can become a bottleneck at scale, making high-bandwidth interconnects needed for maintaining training efficiency. Practical systems often limit expert parallelism to within-node communication to exploit NVLink bandwidth.
- Implementation requires careful buffer management for variable-length messages and coordination with routing decisions to preserve token ordering through the dispatch-compute-combine pipeline.
Expert parallelism combines naturally with other parallelism strategies. Production systems typically use expert parallelism for the MoE layers alongside data parallelism for throughput and potentially tensor parallelism for very large expert dimensions. The choice of which strategies to combine, and in what proportions, depends on the specific model architecture, available hardware, and performance requirements. Getting the combination right is one of the central engineering challenges in large-scale MoE training, and it rewards careful analysis of the communication and computation time breakdowns for the specific configuration at hand.
Quiz
Ready to test your understanding? Take this quick quiz to reinforce what you've learned about expert parallelism in Mixture of Experts models.
Expert Parallelism
Reference
Citation details
Cite or share this article.
Continue with the full handbook
This chapter is part of Language AI Handbook. Use the handbook page to browse the complete table of contents and continue reading in sequence.
Explore Language AI HandbookStay up to date
Get articles, book updates, and news delivered to your inbox.
No spam, unsubscribe anytime.
Join the community
Sign in to remove popups, track your reading progress, and join the discussion.

Comments
No comments yet. Be the first to share your thoughts!