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

原文：https://zhuanlan.zhihu.com/p/710296768

在上一篇文章中，我们介绍了大模型训练中的通信优化技术。

大模型训练：Megatron-Core中的通信优化

除了通信优化，对于大模型训练来说，内存优化也是非常重要的问题。在这篇文章里，我们会结合快手发表在今年ATC上的工作（Megatron-Kwai[1]），简单介绍一下内存优化技术。

背景

首先，根据生命周期的长短，我们将内存开销分为：长内存、中内存、和短内存。

长内存在训练任务开始以后就不会消失，包括模型参数、模型梯度、和优化器参数（例如Adam中的一阶矩估计和二阶矩估计）。中内存会跨越不同阶段，例如由前向传递得到的激活值会保存到后向传递用于计算梯度。短内存只跨越几个算子，或者仅仅在算子计算过程中使用（workspace内存）。

在PyTorch框架中，以上内存开销可以通过memory_allocated查看 ^[2] 。除此以外，我们还有CUDA context、NCCL buffer等开销。在这篇文章中，我们主要分析长内存开销和激活值开销 。

LLaMA模型

以LLaMA模型为例，大模型结构包含一个embedding词表，L个transformer层，以及一个分类头。对于每一层transformer layer，我们统计所有线性层的参数，可以得到模型参数为 ：

内存优化技术

针对激活值开销，在Megatron-Kwai中，主要介绍了两种内存优化手段：重计算和CPU卸载。

首先，重计算技术在前向传递的时候不保存激活值，而在反向传递前重新进行一次前向计算，是一种用时间换空间的手段。在Megatron-LM中，支持全部的重计算和层级的选择性重计算，也就是说，我们可以对所有的transformer layers或者部分的transformer layers进行重计算。

其次，我们可以使用CPU卸载的方法继续降低激活值内存。在流水线并行中，我们将每个model chunk中的激活值卸载到CPU。在反向传递阶段，我们提前一步将CPU上的激活值加载到GPU中。