← 研踪 PaperTrailarXiv:2611.11364 · v1
OOM: Offload-Optimized Full-Parameter MoE Training via Memory-Execution Co-Design
首次提交 2026-11-08 21:16
版本更新 2026-11-08 21:16
雷达观测 2026-11-10 12:16
摘要
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.
作者备注：19 pages, 9 figures. Code: https://github.com/RegiaYoung/OOM-MoE
观测到的版本
相邻记录可能跳过中间版本；首次发现 v2 及以上也会标为版本更新。
v1 · 2026-11-08 21:16 · 最新观测版本
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.
查看此版本解读 →查看此版本原文 ↗