Expert Parallelism › Expert Placement Strategies
Language AI HandbookAll chapters Expert Parallelism › Expert Placement StrategiesOne section · Sections in this chapter3 / 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
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.