MemFine: Memory-Aware Fine-Grained Scheduling for MoE Training

2025-11-26

AI Summary

Large-scale Mixture-of-Experts (MoE) model training suffers from severe load imbalance induced by dynamic token routing, causing GPU memory overflow and hindering model scalability. Existing capacity-constrained load-balancing approaches sacrifice accuracy and fail under memory-constrained hardware. This paper proposes a fine-grained scheduling framework: first decoupling token distribution from expert computation units; then establishing a theory-driven memory model; and finally designing a dynamic optimization strategy for recomputation chunking, integrated with memory-aware scheduling and real-time load balancing. The method enables memory-efficient and stable training without compromising model accuracy. Experiments demonstrate a 48.03% reduction in activation memory, a 4.42× throughput improvement, and—critically—the first reliable training of ultra-large MoE models on low-memory GPUs.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Planning, Routing, and Scheduling: Planning with Language Models
• Search and Optimization: Learning to Search

Application Category

• User Modeling, Personalization and Recommendation: Practical large-scale studies of user experience
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

The training of large-scale Mixture of Experts (MoE) models faces a critical memory bottleneck due to severe load imbalance caused by dynamic token routing. This imbalance leads to memory overflow on GPUs with limited capacity, constraining model scalability. Existing load balancing methods, which cap expert capacity, compromise model accuracy and fail on memory-constrained hardware. To address this, we propose MemFine, a memory-aware fine-grained scheduling framework for MoE training. MemFine decomposes the token distribution and expert computation into manageable chunks and employs a chunked recomputation strategy, dynamically optimized through a theoretical memory model to balance memory efficiency and throughput. Experiments demonstrate that MemFine reduces activation memory by 48.03% and improves throughput by 4.42% compared to full recomputation-based baselines, enabling stable large-scale MoE training on memory-limited GPUs.

Problem

• Addresses memory bottleneck in MoE training from load imbalance
• Solves GPU memory overflow issues constraining model scalability
• Overcomes accuracy compromises of existing capacity-capping methods

Innovation

• Decomposes token distribution into manageable chunks
• Uses chunked recomputation strategy for memory efficiency
• Employs theoretical memory model for dynamic optimization