Skip to content
AI360Xpert
Core ML

Expert Parallelism

When a Mixture of Experts model has too many experts to fit on one GPU, you distribute the experts across multiple GPUs. But now, every token must be sent across the network to its assigned expert.

In Expert Parallelism, tokens are routed across the network to different GPUs. GPU 1 might hold Expert A, while GPU 2 holds Expert B. The router decides where each token goes.
In Expert Parallelism, tokens are routed across the network to different GPUs. GPU 1 might hold Expert A, while GPU 2 holds Expert B. The router decides where each token goes.

The MoE Scaling Problem

Mixture of Experts (MoE) models scale up parameter counts without increasing compute costs. A model like Mixtral 8x7B has 47 billion parameters, but only uses about 13 billion parameters per token.

However, memory is still a bottleneck. Even if a token only uses 13 billion parameters, all 47 billion parameters must be loaded into GPU memory. For larger MoE models (like GPT-4's rumoured architecture with over a trillion parameters), no single node can hold the entire model.

Enter Expert Parallelism

Expert Parallelism (EP) is a distributed training strategy specifically designed for MoE models.

Instead of sharding every layer across all GPUs (like tensor-parallelism), EP assigns whole experts to specific GPUs.

  • GPU 0 stores Expert 1 and Expert 2.
  • GPU 1 stores Expert 3 and Expert 4.
  • GPU 2 stores Expert 5 and Expert 6.
  • GPU 3 stores Expert 7 and Expert 8.

The non-MoE layers (like self-attention) are typically parallelized using standard Data Parallelism or Tensor Parallelism.

The All-to-All Bottleneck

Expert Parallelism introduces a massive network communication challenge: All-to-All communication.

In a standard batch of text, the Router assigns different tokens to different experts. Because the experts now live on different GPUs, the tokens themselves must be sent over the network to the correct GPU, processed, and then sent back.

  1. Routing: The Router looks at a batch of tokens and decides which expert gets which token.
  2. Scatter (All-to-All): Every GPU sends its tokens to the GPUs that hold the corresponding experts. This is an all-to-all operation because every GPU is sending data to every other GPU simultaneously.
  3. Compute: The experts process the tokens they received.
  4. Gather (All-to-All): The GPUs send the processed tokens back to the GPUs they originally came from, reassembling the batch.

If the network interconnect (like NVLink or InfiniBand) is slow, the GPUs will spend more time waiting for tokens to arrive than they will spend actually computing.

Load Balancing and Dropped Tokens

Another major challenge in Expert Parallelism is Load Balancing.

If the Router decides that 90% of the tokens in a batch should go to Expert 1, then GPU 0 (which holds Expert 1) will be overwhelmed, while GPUs 1, 2, and 3 sit idle. This destroys training efficiency.

To solve this, researchers use capacity limits and auxiliary loss functions:

  • Capacity Limits: An expert is only allowed to process a maximum number of tokens per batch (e.g., 100 tokens). If 150 tokens are routed to it, the extra 50 tokens are "dropped" and passed to the next layer unchanged.
  • Auxiliary Loss: A penalty is added to the training loss if the Router does not distribute tokens evenly across all experts.

The Quick Version

  • Expert Parallelism distributes the experts of an MoE model across multiple GPUs to save memory.
  • It requires All-to-All communication to send tokens over the network to their assigned experts and back.
  • It relies on fast interconnects (like NVLink) to prevent the network from becoming a bottleneck.
  • Load balancing is critical; if the router favours one expert, that GPU becomes a bottleneck while others sit idle.

Related concepts