Zyphra Cuts MoE Training Communication Costs by 2.63x Using Router Patterns

At scale, Mixture-of-Experts training can spend more time moving tokens than running expert computation. A paper from Zyphra reports that router choices become predictable early in pretraining, allowing the runtime to co-locate frequently paired experts and move tokens toward likely destinations.

Benchmarks in the paper, run with Megatron-LM on AMD Instinct MI300X clusters, accelerated all-to-all exchanges by 1.16× to 2.63× and the full training step by up to 1.41×.

MoE’s network bottleneck

In an MoE layer, a router assigns each token to its top-k experts. Expert parallelism, or EP, distributes those experts across GPU ranks; EP32 means the expert group spans 32 GPUs. Each forward pass dispatches tokens to the selected experts and returns the resulting activations, while backpropagation requires corresponding data movement.

On clusters with eight MI300X GPUs per node, these all-to-all collectives consume about 45% of an EP32 training step with top-2 routing and 60% with top-6 routing. Higher values of k create more destinations per token, and EP groups that cross node boundaries encounter slower interconnects.

Existing MoE systems reduce this cost through communication overlap, replicated experts, or modified routing schemes. Zyphra’s design operates in the placement and dispatch layers, using observed routing patterns to reduce the amount of data crossing the network.