OOM: Offload-Optimized Full-Parameter MoE Training via Memory-Execution Co-Design

원문 초록

Offloading model state to host memory and streaming layer weights let a single multi-GPU node train large mixture-of-experts (MoE) models. GPU memory then becomes a performance budget, shared by resident expert state, retained activations, and execution worksets. The best allocation depends on the expert execution path, batch size, and sequence length. We present OOM (Offload-Optimized MoE), a full-parameter MoE training system that co-designs expert execution and memory allocation. For expert-parallel execution, OOM updates resident experts on their owning GPUs and jointly plans state residency and MoE activation retention according to pipeline dependencies. For local execution, OOM bounds GPU worksets through dependency-aware partitioning; in hybrid models, it also retains attention results on the host to avoid costly recomputation and selectively regenerates layer boundaries through cheaper layers to make room. Calibrated cost models guide joint configuration within GPU and host budgets to minimize step time; OOM selects among the configured paths through short executions. On four H100s, OOM trains with 1M-token contexts, whereas FSDP2, ZeRO-Offload, and MegaTrain exhaust GPU memory at 512K. It also trains a 122B-parameter MoE model, on which all three exhaust host memory. On four 24 GB RTX 4090s, it trains 512K-token contexts and 1,024 1K-token sequences per step without gradient accumulation. Across matched-global-batch comparisons, OOM delivers 1.58x and 4.08x geometric-mean (up to 2.14x and 7.77x) throughput over CPU-offloaded FSDP2 on the two platforms while using the least host memory among the evaluated systems -- as little as 1/3.0 of FSDP2's.