Reducing Activation Recomputation in Large Transformer Models

摘要

本文展示了如何通过减少激活重计算来提升训练吞吐，激活重计算通常用于解决内存限制，其将激活值丢弃，再在需要的时候重新计算，以计算换内存。在这项工作中，证明了大部分冗余计算是不必要的，因为有其他方法可以充分减少内存消耗，并且不会带来额外成本。提出了两个优化手段：序列并行和选择激活重计算。

在万亿规模的模型上评估了我们的方法，结果表明优化方法将激活内存减少了5倍，同时将激活重计算的执行时间开销减少了90%以上。具体的，在 2240 张 A100 上训练 530B 的模型时，能够达到 54.2% 的 MFU（Model Flops Utilization），这比使用激活重计算进行训练达到的 42.1% 快了 29%

1. 引言

模型并行技术能支持训练更大规模的模型，但也存在一些限制。

张量并行增加了通信需求，并引入了更小且性能较差的矩阵乘运算，这导致在大量设备上切分模型后训练效率降低。因此，张量并行一般在节点内进行切分。

流水线并行需要存储多个 microbatch 的激活值以减小气泡。因此，流水线并行只能 offload 模型参数和优化器状态，而无法在保持高设备利用率的情况下减少激活内存占用。

<aside> 💡 插个题外话：如何减小气泡？

增加 microbatch 数，可以减小气泡，关系公式：$\frac{p-1}{m}$；
切分更细粒度的 pipeline stage，为每个设备分配多个(e.g. 2) stage，进行交错执行，也可以减小气泡时间（2倍），但是会增加通信量（2倍）；

</aside>

下图量化了不同参数规模的模型所需要的内存大小：

Untitled

可以看出，就算是 22B 的模型，所需要的内存也超过了单卡的显存容量（80GB）。目前常用的解决方法是将激活值舍弃，并在需要时重新计算以在反向过程中计算梯度，该方法称为：“梯度检查点” 或者 ”激活重计算“。但这种方法会带来额外的计算开销，降低训练效率，这种方法的冗余计算所带来的时间开销达到了 30%~40%。

下面会对存储 Transformer 模型的激活值所需要的内存大小建立一个近似公式，通过这个公式，我们可以分析不同形式的模型并行对激活内存大小的影响。