Mixture-of-Parallelisms: Towards Memory-Efficient Training Stack for Mixture-of-Experts Models

Jul 16, 2026

This paper proposes a memory-efficient training stack for Mixture-of-Experts (MoE) models, combining existing and new parallelism techniques across different layers and stages of the MoE training pipeline. The approach targets maximal efficiency under practical CPU/GPU and interconnect bandwidth constraints, and introduces an optimizer-step strategy to improve throughput and memory usage. Experiments show 4.7–8.2× higher per-GPU throughput than a tuned FSDP2 baseline and enable training with up to 1M tokens, avoiding out-of-memory issues beyond 64–128K for the baseline.

LanguageEnglish

CategoryResearch & Report

FormatPDF

Pages9

Reading Time23 min

Introduction

Limitation of one-size-fits-all parallelism

Mixture-of-Parallelisms (MoP) approach

Lossless principles for trillion-parameter, million-token scale

arXiv :2607 .0 1844v 1 [ cs .DC] 2 Jul 2026 Mixture-of-Parallelisms: Towards Memory-Efficient Training Stack for Mixture-of-Experts Models Xuan-Phi Nguyen†∗ Shrey Pandit† Yiran Zhao† Semih Yavuz Silvio Savarese Shafiq Joty∗ Salesforce AI Research Abstract This paper proposes a memory-efficient training stack for Mixture-of-Experts (MoE) models. It is a training paradigm that combines and specializes various existing and novel parallelism techniques at different layers and stages of the Mixture-of-Experts (MoE) model training pipeline. It leverages these techniques to achieve maximal efficiency given the physical constraints of CPU, CPU memory, GPU HBM memory, and the CPU-GPU, GPU-GPU, and node-node communication bandwidth of the GPU cluster. It also contains a novel strategy for the optimizer step to achieve high throughput and memory efficiency, enabling practitioners to conduct lossless pre-training/fine-tuning of trillion-parameter scale models, at a million context length, with just under 12 8x H200 GPU nodes, with state-of-the-art throughput and memory efficiency. In our experiments, MoP delivers 4.7–8.2 × higher per-GPU throughput than a strongly-tuned FSDP2 baseline (with the gap widening at larger scale) and sustains training at context lengths up to 1M tokens, where the baseline runs out of memory beyond 64–128K. 1 Introduction Sparse Mixture-of-Experts (MoE) architectures have become the dominant recipe for scaling language models, because they decouple parameter count from per-token compute: routing each token to only k of E experts lets the parameter budget—and with it model quality—grow far faster than the FLOPs per token [6, 3] . The frontier this enables, trillion-parameter models trained at near-million-token context, is no longer compute-bound on modern accelerators; it is bound by GPU memory and by the communication bandwidth of the cluster. A single trillion-parameter model needs ∼ 16Θ bytes of persistent optimizer-related state alone [8](§2.1), tens of terabytes that no single device can hold, while long-context activations add a transient footprint that grows with sequence length and depth. Training such models therefore hinges on how effectively a system spreads state and traffic across the memory hierarchy and interconnects of the machine. The limitation of one-size-fits-all parallelism. The prevailing systems attack this problem with a single global plan applied uniformly to every layer. Megatron-LM [13, 5] composes tensor, pipeline, and sequence parallelism into a 3D grid W = dtp; ZeRO and its PyTorch realization FSDP [10, 14] shard the persistent state across all data-parallel ranks; expert parallelism [11] partitions experts across devices. Each is excellent at the bottleneck it targets and suboptimal elsewhere. Yet the pressures of trillion-scale long-context MoE training are sharply component-dependent: attention is activation-bound, the MoE feed-forward is weight-and routing-bound, the vocabulary projection produces the single largest logit tensor, and the optimizer state dwarfs the weights themselves. Noone parallelism is Pareto-optimal across all of these—ZeRO-3 shards the expert weights but leaves ∗ Project lead & corresponding authors: {xnguyen,[sjoty}@salesforce.com](sjoty}@salesforce.com) System Card. Preliminary Work In Progress. the full-rank MoE activations on every device, expert parallelism shards experts but not the O(NkH) routed activations, and 3D parallelism multiplies its axes against a fixed device budget while still applying the same trade-off to every layer. The result is wasted headroom: at any moment most of the machine’s resources—idle CPU cores, host DRAM, and spare interconnect bandwidth—sit unused while one bottleneck dictates the global plan. Mixture-of-Parallelisms. We propose Mixture-of-Parallelisms (MoP), a training stack that abandons the single global plan and instead specializes a different parallelism to each component of the MoE block, ch

FAQs

What problem does the Mixture-of-Parallelisms (MoP) training stack address?

MoP addresses the inefficiency of using a single global parallelism plan for all MoE components during trillion-parameter, long-context training, which wastes CPU, host memory, and interconnect headroom when one bottleneck dominates.

How does MoP differ from existing parallelism strategies like Megatron-LM and ZeRO/FSDP?

Instead of applying one uniform parallelism trade-off across every layer, MoP assigns different parallelism methods to different MoE components based on the component’s dominant bottleneck, reducing wasted resources and rank multiplication effects.

What performance and scalability results are reported in the paper?

The experiments report 4.7–8.2× higher per-GPU throughput than a strongly-tuned FSDP2 baseline, with the advantage widening at larger scale, and training sustained up to 1M tokens where the baseline runs out of memory beyond 64–128K.

Jul 16, 2026

Upload by Olivia Brown