Expert Parallelism › All-to-All Communication
Language AI HandbookAll chapters Expert Parallelism › All-to-All CommunicationOne section · Sections in this chapter4 / 12 ↓
Reading tools Definitions & inline links
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 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.