Expert Parallelism › Expert Parallelism Implementation
Language AI HandbookAll chapters Expert Parallelism › Expert Parallelism ImplementationOne section · Sections in this chapter6 / 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
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.