大模型训练：Megatron-Kwai中的内存优化

文章预览

原文：https://zhuanlan.zhihu.com/p/710296768 在上一篇文章中，我们介绍了大模型训练中的通信优化技术。 大模型训练：Megatron-Core中的通信优化 除了通信优化，对于大模型训练来说，内存优化也是非常重要的问题。在这篇文章里，我们会结合快手发表在今年ATC上的工作（Megatron-Kwai[1]），简单介绍一下内存优化技术。 背景 首先，根据生命周期的长短，我们将内存开销分为：长内存、中内存、和短内存。 长内存在训练任务开始以后就不会消失，包括模型参数、模型梯度、和优化器参数（例如Adam中的一阶矩估计和二阶矩估计）。中内存会跨越不同阶段，例如由前向传递得到的激活值会保存到后向传递用于计算梯度。短内存只跨越几个算子，或者仅仅在算子计算过程中使用（workspace内存）。 在PyTorch框架中，以上内存开销可以通过memory_allocated查看 [2] 。除此以 ………………………………