消除长上下文混合专家训练中的每个内存峰值
研究人员提出针对长上下文混合专家（MoE）模型训练的内存优化技术。以往只要任一组件峰值内存超过设备内存，训练就会失败。研究识别出四个瓶颈：专家分发、词汇投影、梯度检查点边界和优化器状态。PipelinedLLEP、Ring-DTP、选择性检查点卸载（SCO）和 OffloadStreamAdamW 以固定 GPU 工作集限制这四者。这些方法仅改变计算顺序和粒度，因此损失和梯度保持精确。测试显示，MoE 分发峰值最多降低 59.3%，词汇投影峰值降低 86.6%，卸载优化器步骤加快 2.05 倍。在 120B 至 667B 参数的 MoE 模型上，他们以 1M 上下文长度训练，达到调优后 FSDP2 基线的 8 至 32 倍范围，吞吐量最高为其 10.4 倍。
本文来源：Hugging Face，仅供学习参考，版权归原作者所有。