SeqMoE improves memory offloading for faster AI model use
SeqMoE: Toward Full-Load Performance via Predictive and Graph-Compatible MoE Offloading
Operating SystemsArtificial Intelligence
Summary
Large AI models often need to save parts of their memory to slower storage, which slows down how fast they work. The authors created SeqMoE, a system that predicts which parts of the model will be needed soon and loads them early to avoid delays. It also uses clever scheduling and caching strategies to keep the right data ready, improving speed without having to keep everything in fast memory. This makes large AI models run much closer to their maximum speed even with limited fast memory.
What this means in practice
- •For machine learning engineers: Improve deployment of large MoE models on limited-memory devices by reducing slow data loads through predictive expert offloading.
- •For high performance computing teams: Design runtime systems that handle expert data offloading efficiently to maximize throughput in large-scale AI workloads.
Authors
Zihan Wang, Yuqi Wang, Lei Gong, Cheng Tang, Wenqi Lou, Teng Wang, Chao Wang, Xuehai Zhou
Abstract
Mixture-of-Experts (MoE) creates a structural advantage for offloading: only a small fraction of activated experts need to reside in device memory, and if they can be loaded in time for computation, offloading can in principle approach full-load performance, where all model weights reside in device memory. Yet translating MoE's structural advantage into practical offloading gains remains challenging. We propose SeqMoE to bridge this gap. To maximize expert hits, we build predictive memory management: (i) Sequence-to-sequence prediction. We are the first to recast expert activation prediction as sequence modeling, enabling accurate multi-step, multi-layer forecasts that provide a long and reliable window for downstream decisions. (ii) Joint prefetch scheduling. We formulate prefetch scheduling as Job Sequencing with Deadlines to maximize expected expert hits and improve bandwidth efficiency. (iii) Forecast-driven caching. Leveraging the recursive nature of sequence modeling, we introduce a probabilistic Belady policy for future-aware eviction. To eliminate execution bottleneck, we develop (iv) Graph-compatible offloading runtime. We derive general runtime principles encompassing compute-transparent expert placement and synchronization-free orchestration disciplines for end-to-end graph capture. With 45% expert residency, SeqMoE averages a 96.97% hit rate and 80.22% of full-load performance, advancing the state of the art in MoE offloading.