SYMI: Efficient Mixture-of-Experts Training via Model and Optimizer State Decoupling

TL;DR

SYMI addresses the convergence and efficiency challenges of training large Mixture-of-Experts models by decoupling the placement of expert parameters from their optimizer state, enabling per-iteration adaptive replication with minimal overhead. It introduces a holistic design with a static, host-resident optimizer state and dynamic, popularity-driven expert replication, coordinated by a dedicated Expert Placement Scheduler and a lightweight Symi Optimizer. New collective communications and placement materialization techniques ensure data movement remains invariant while permitting fine-grained rebalancing, yielding significant speedups over state-of-the-art baselines (DeepSpeed and FlexMoE). The approach demonstrates robust improvements in time-to-convergence and token efficiency on multiple GPT-scale models, highlighting practical benefits for training massive MoE models without prohibitive rebalancing costs.

Abstract

Mixture-of-Experts (MoE) models have become a widely-adopted solution to continue scaling model sizes without a corresponding linear increase in compute. During MoE model training, each input token is dynamically routed to a subset of experts -- sparsely-activated feed-forward networks -- within each transformer layer. The distribution of tokens assigned to each expert varies widely and rapidly over the course of training. To handle the wide load imbalance across experts, current systems are forced to either drop tokens assigned to popular experts, degrading convergence, or frequently rebalance resources allocated to each expert based on popularity, incurring high state migration overheads. To break this performance-accuracy tradeoff, we introduce SYMI, an adaptive MoE training system. The key insight of SYMI is to decouple the placement of expert parameters from their large optimizer state. SYMI statically partitions the optimizer of each expert across all training nodes. Meanwhile, SYMI dynamically adjusts the placement of expert parameters by repurposing existing weight updates, avoiding migration overheads. In doing so, SYMI right-sizes the GPU resources allocated to each expert, on a per-iteration basis, with minimal overhead. Compared to state-of-the-art MoE training systems, DeepSpeed and FlexMoE, SYMI is able to achieve a 30.5% and 25.9% faster time-to-convergence, respectively.

Figures (13)

• Figure 1: Overview of an MoE layer and expert parallelism.
• Figure 2: A single layer's expert popularity distribution during the training of GPT-Small (125M) extended with 32 experts. The distribution shifts dramatically within very few iterations.
• Figure 3: Symi decouples the model and optimizer state placement. Expert replication is non-uniform and dynamic, while each expert's optimizer remains static, uniformly sharded across all hosts.
• Figure 4: Symi design block diagram. Showing data (solid lines) and metadata (dashed lines) flow in a single rank through a training iteration.
• Figure 5: Current systems (left) bind optimizer state to expert instances, requiring costly optimizer state migration during expert rebalancing. Symi (right) keeps the optimizer static and eliminates this overhead.