Expert Parallelism: The Router's Dilemma

We have explored Tensor Parallelism (TP), which slices individual operations to fit massive layers into memory, and Pipeline Parallelism (PP), which slices layers vertically to scale across nodes. But as models grow not just in depth or width, but in breadth—specifically with Mixture-of-Experts (MoE) architectures like Mixtral 8x7B or DeepSeek-V3—we face a new bottleneck.

In a dense model, every parameter is used for every token. In an MoE model, only a fraction of parameters (experts) are active per token. Replicating every expert on every GPU can be too expensive in memory. Expert parallelism shards the experts across ranks; it can also be combined with tensor and data parallelism. The best layout depends on expert size, batch size and interconnect.

Enter Expert Parallelism (EP). In this paradigm, we distribute the experts themselves across different GPUs. GPU 0 holds Expert A, GPU 1 holds Expert B. The challenge shifts from splitting matrices to routing tokens: we must physically move data (tokens) to the GPU that holds the parameters required to process them.

This article dissects the hardware mechanics of EP, the "All-to-All" communication primitive that defines it, and the critical trade-offs involved.

1. The Hardware Mechanics: Token Routing

In Expert Parallelism, the model is physically fragmented. The critical hardware event is no longer a matrix multiplication synchronization (like TP) or a pipeline flush (like PP); it is token movement across the interconnect. The volume depends on the number of routed tokens, hidden width and expert placement.

The Four-Step Lifecycle

For a single MoE layer, the hardware execution flow looks like this:

1. Route (Local Compute): Each GPU processes its batch of tokens through a small "Gating Network" (usually a single Linear layer). This router calculates which expert (e.g., Expert 0, 1, ... 7) each token needs to visit.

   • Hardware state: Tokens are currently residing on the "source" GPU.

2. Dispatch (The Communication Bottleneck): This is the defining moment of EP. GPU 0 might have 100 tokens that need to go to Expert 5 (on GPU 5). GPU 1 might have 50 tokens for Expert 5. Every GPU must send specific subsets of its data to every other GPU simultaneously. A common implementation uses an All-to-All collective operation.

   • Hardware implication: This stresses the bisection bandwidth of the cluster. All-to-All has a different traffic pattern from the collectives often used for TP.

3. Expert Compute (Local Compute): Once the tokens arrive at their destination GPUs (the "home" of their assigned experts), the GPU performs the standard Feed-Forward Network (FFN) computation.

   • Hardware benefit: Arithmetic intensity depends on the number of tokens assigned to each expert and the expert dimensions. The expert computation itself can run locally.

4. Combine (Communication): The processed tokens must return to their original source GPUs to preserve the sequence order for the next layer (e.g., Attention). This requires a second All-to-All operation to send results back.

The Hardware Reality: All-to-All

The All-to-All primitive is distinct from All-Reduce.

• All-Reduce (used in some TP/DP operations): Ranks combine equal-shaped tensors with a reduction operation, such as summation. The communication volume depends on tensor size and algorithm.
• All-to-All (EP): GPU $i$ sends a personalized buffer to GPU $j$. The volume can be dynamic and uneven (load imbalance), causing network congestion if one GPU receives significantly more data than others.

2. Bare Metal Implementation: Pure PyTorch

This top-1, one-expert-per-rank pseudocode shows the routing order. PyTorch torch.distributed.all_to_all_single accepts flat tensors and integer split-size lists for the two variable-size transfers. A real layer also needs consistent router weights, expert output weighting, a process group, gradient handling and any capacity policy.

# Schematic top-1 dispatch; one expert per rank. Not executable code.
logits = router(local_tokens)           # router weights identical on all ranks
expert = argmax(logits, axis=-1)         # one destination rank per token
order = argsort(expert)                  # group tokens by destination
send_counts = bincount(expert, minlength=world_size)
recv_counts = all_to_all(send_counts)    # exchange counts before variable splits
received = all_to_all(sorted_tokens, input_splits=send_counts,
                      output_splits=recv_counts)
processed = local_expert(received)
returned = all_to_all(processed, input_splits=recv_counts,
                      output_splits=send_counts)
output = inverse_permute(returned, order).reshape_as(local_tokens)

The sketch shows data movement, not a runnable PyTorch layer. Megatron Core supports several dispatcher and token-capacity strategies.

In this implementation, the GPUs effectively swap roles: they stop being "owners of a batch slice" and temporarily become "compute servers for specific operations," processing tokens from anywhere in the cluster.

3. Hardware Constraints & Trade-offs

The All-to-All Latency Spike

EP moves routed activations twice per MoE layer in this example. The volume is proportional to routed token count and hidden width. TP and PP also communicate, and any of these strategies may run within or across nodes.

• Inter-Node Bottleneck: If your experts are spread across different server nodes, the All-to-All operation must traverse the Ethernet/InfiniBand network. Cross-node bandwidth and latency generally differ from local links; measure the actual topology.
• Latency Sensitivity: For small batch sizes (inference), the latency of the All-to-All step can dominate the execution time, making EP slower than simple data parallelism.

The Load Balancing Nightmare

Hardware hates irregularity. GPUs are fastest when processing static, uniform shapes.

• The Problem: If the Router decides that 90% of the tokens need to go to Expert 0 (GPU 0), then GPU 0 becomes a bottleneck while GPUs 1–7 sit idle. This is computational skew.
• The Consequence: The overloaded expert can delay the collective and subsequent computation.
• The Solution (Hardware Level): Possible remedies include routing/load-balancing methods, expert replication, larger capacity or token dropping. Dropping is an optional capacity policy and changes model behavior.

Pros and Cons Summary

Pros

• Massive Parameter Scale: You can train models with trillions of parameters — Google's Switch Transformer reached 1.6T parameters with 2,048 experts — because expert weights can be sharded across devices. The $1/N$ fraction applies only to the expert weights in the idealized one-expert-per-rank layout; shared attention and other weights remain.
• Compute Efficiency: Unlike dense models where FLOPs scale linearly with parameter count, For a fixed number and size of active experts, adding inactive experts can increase total parameter count without proportionally increasing expert FLOPs per token. Router, communication and memory costs still matter; quality gains are empirical.
• Asynchronous Potential: Some implementations overlap expert communication with computation when dependencies permit.

Cons

• Communication Overhead: The All-to-All primitive is heavy. It benefits from sufficient bandwidth in the participating process group.
• Load Balancing: Routers need a balancing strategy; auxiliary losses are one option. Skew can reduce utilization.
• Memory Fragmentation: Dynamic token counts can require variable-size buffers or padding; the chosen dispatcher determines memory overhead.

Sources

• PyTorch distributed all_to_all_single API
• NVIDIA Megatron Core MoE features and dispatchers
• Switch Transformers paper

Related

• Tensor Parallelism Deep Dive — sharding the weights themselves.
• Pipeline Parallelism Deep Dive — the bubble problem when you shard by layer instead.