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

本文来源公众号“GiantPandaCV”，仅用于学术分享，侵权删，干货满满。

原文链接：大模型训练：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中。

在这个过程中，激活值的加载/卸载可以和前向/反向计算overlap。Megatron-Kwai允许调整卸载的比例（offload_ratio），选择将部分激活值卸载到CPU上。如果数据卸载的时间小于计算的时间，我们可以将这部分开销全部隐藏。

总结

对于内存优化，主要的手段包括并行技术、重计算、以及CPU卸载。一方面，我们使用张量并行和流水线并行来降低模型参数对应的长内存开销。另一方面，我们使用序列并行和文本并行来降低模型输入对应的激活值开销。最后，合理使用重计算和CPU卸载技术，可以用很小的成本来降低激活值内存，这有助于我们选择更加高效的并行配置来训练大模型。

THE END !

文章结束，感谢阅读。您的点赞，收藏，评论是我继续更新的动力。大家有推荐的公众号可以评论区留言，共同学习，一起进步。