论文推荐|固定GPU内存搞定大规模MoE长上下文训练:1M上下文吞吐最高提升10.4倍
thinkindev • 2026-09-23
1535 views
一项针对大规模混合专家(MoE)模型长上下文训练内存瓶颈的研究提出四项调度技术,目标是同时约束所有内存峰值,而非仅降低平均占用。研究指出,长上下文或大batch下,任一组件的峰值分配超过设备内存就会失败,因此必须对专家分发、词表投影、梯度检查点边界和优化器状态分别有界化。PipelinedLLEP在最少负载专家并行中限制每个源对分发块的token贡献;Ring-DTP在词表投影处通过环形流通激活或权重分片,并把logits块折叠到在线log-sum-exp;选择性检查点卸载(SCO)将每个检查点边界的唯一长寿命张量保留在CPU内存;OffloadStreamAdamW则把串行CPU Adam更新改为桶式流水线。四项技术只改变计算与数据移动的顺序和粒度,损失和梯度保持精确。在组件测试中,MoE分发峰值最高降低59.3%,词表投影峰值降低86.6%,卸载优化器步骤加速2.05倍;在120B至667B参数MoE模型上组合应用后,可训练1M上下文长度,可及范围是调优FSDP2基线的8至32倍,吞吐最高达10.4倍。
核心要点
- 四项调度技术分别约束专家分发、词表投影、检查点边界和优化器状态四大内存峰值
- 组件测试显示MoE分发峰值最高降59.3%、词表投影峰值降86.6%,优化器步骤加速2.05倍且精度无损
- 在120B至667B MoE模型上实现1M上下文训练,可及范围是FSDP2基线8-32倍,吞吐最高10.4倍