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

Abstract: This paper showcases 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.7x--8.2x 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.

• The paper introduces a Mixture-of-Parallelisms (MoP) strategy that assigns specialized parallelism per MoE component, alleviating memory and bandwidth bottlenecks.
• MoP achieves up to 8.2× higher per-GPU throughput and supports million-token context lengths, enabling practical training on commodity GPU clusters.
• It combines sequence parallelism, least-loaded expert routing, and asynchronous optimizer host-offload to efficiently overlap communication and computation.

Mixture-of-Parallelisms: Memory-Efficient Training for Trillion-Scale MoE Models

Introduction and Motivation

The paper "Mixture-of-Parallelisms: Towards Memory-Efficient Training Stack for Mixture-of-Experts Models" (2607.01844) introduces a new parallelization paradigm tailored for training extremely large Mixture-of-Experts (MoE) LLMs, with a focus on memory and throughput efficiency. While prior systems, such as Megatron-LM and ZeRO/FSDP, use a single global parallelism configuration applied uniformly across all neural network layers, MoE architectures at trillion-parameter and million-token context scales present diverse component-level bottlenecks—across persistent model weights, transient activations, optimizer state, routing, and output projections—that no single parallelism scheme optimally addresses. The Mixture-of-Parallelisms (MoP) stack instead assigns specialized parallelism strategies on a per-component basis, mitigating bottlenecks on memory and communication bandwidth.

The Mixture-of-Parallelisms Architecture

MoP’s core principle is deploying distinct forms of parallelism for each functional sub-component of the MoE transformer block, rather than utilizing a Cartesian composition of pipeline, data, or tensor parallelism. Each parallelism operates over overlapping subsets of $W$ global data-parallel ranks, avoiding exponential rank multiplication.

Figure 1: Overview of the Mixture-of-Parallelisms approach; attention, MoE experts, FFN, and projection each use specialized parallelism on the same $W$ data-parallel group, avoiding Cartesian rank explosion.

Specifically:

• Attention activations employ sequence parallelism, reducing per-rank activation from $\mathcal{O}(SH)$ to $\mathcal{O}(NH)$, where $N=S/D$ and $D$ is the sequence parallelism degree. This allocation is orthogonal to expert sharding, enabling independent scaling.
• MoE expert parameters use a hybrid of expert parallelism and parameter sharding, distributing the dominant parameter mass to minimize per-rank persistent state. Routing volume is decoupled and addressed via Least-Loaded Expert Parallelism (LLEP).
• Feed-Forward Network (FFN) and dense path parameters are sharded efficiently across $W$ ranks, leveraging their small relative parameter footprint.
• Vocabulary projection is handled with a data-tensor parallel scheme that avoids fully materializing the $N \times V$ logit matrix, providing loss and gradient computation with only a fraction of memory.
• Optimizer state (AdamW’s master weights and moments) is offloaded to host memory and updated via a pipelined asynchronous mechanism that overlaps communication and computation, maintaining a small device working set.

This per-component parallelism avoids the strict factorization of ($d, t, p$) in 3D parallelism, thereby flexibly matching the hardware—GPU HBM, host DRAM, CPU/GPU, and inter-node bandwidths—to workload bottlenecks.

Technical Innovations

Memory-Efficient Expert Parallelism

MoP leverages LLEP, which dynamically migrates expert computation to balance device load according to observed token-expert traffic. An additional memory-efficient variant overlaps the routing communication with the expert computation so that the transient MoE activation peak is lowered—activations, dispatch, and outputs are not all resident in memory simultaneously. This enables scaling token/batch size at fixed hardware resources without compromising speed.

Sharded Data-Tensor-Parallel Vocabulary Projection

Vocabulary projections are sharded column-wise, but unlike standard tensor parallelism, each rank retains a unique batch (DP-like property). The exact loss and gradients are computed through per-block partial aggregation, requiring only a partial $N \times V/P$ activation resident per rank at any time. The approach yields $W$0-$W$1 reduction in memory for the projection step at the expense of negligible context-length-independent communication.

Asynchronous Optimizer Host-Offload

Given that optimizer state for trillion-parameter models is an order of magnitude larger than the weight parameters, device memory cannot accommodate them. MoP maintains optimizer state on host memory and utilizes a pipelined update scheme that overlaps communication with backward propagation, essentially removing optimizer update latency from the critical path. This additionally improves throughput, even at the largest model and context scales.

Empirical Results

Comprehensive benchmarks show MoP achieves $W$2–$W$3 higher per-GPU throughput than the best-tuned FSDP2 baseline, especially as model scale and context length increase. Notably, MoP sustains training up to 1M-token context lengths—where baselines fail at $W$4–$W$5K tokens due to memory exhaustion—while throughput remains nearly constant even as context grows.

The results are summarized in the table below:

Model Size	Nodes	Max Context FSDP2 Fits	MoP Speedup
120B	2 × 8 H200	128K	4.7×
600B	8 × 8 H200	64K	6.1×
1T	12 × 8 H200	128K	6.6–8.2×

These findings demonstrate that MoP enables trillion-parameter, million-context training with a practical twelve-node cluster, a regime previously inaccessible except on much larger hardware budgets.

Limitations and Future Directions

The principal trade-off in MoP is increased communication complexity: parameter sharding and expert routing require frequent collective and all-to-all operations, making MoP communication-bound in bandwidth-limited scenarios. Performance deteriorates if network interconnect is oversubscribed or at very large device counts. In such settings, standard pipelined or replicated data parallelism may partially recover throughput at greater memory cost.

Automating MoP’s sub-group sizing (for sequence, expert, and projection parallelism) based on model structure and physical topology is an important future direction, as is extending the component-specialized scheme to large-context inference and heterogeneous system architectures.

Conclusion

Mixture-of-Parallelisms presents a component-specialized sharding paradigm for MoE models, replacing the uniformly applied parallelism plans of prior art with bottleneck-sensitive assignments for each architectural component. Its lossless overlap of communication and computation, memory-adaptive expert routing, and sharded projection layers jointly remove the major capacity constraints of prior approaches. Empirically, MoP makes trillion-parameter, million-token-context training practical for commodity-scale GPU clusters, suggesting that future scaling of sparse models will profit most from flexible, component-level parallel strategies rather than monolithic global ones. Continued progression in this direction will likely prompt new automated scheduling algorithms and tailored system-architecture codesign to further efficiency and scale.