AI学习者 · 8月12日

大模型训练: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,我们统计所有线性层的参数,可以得到模型参数为:

image.png

内存优化技术

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

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

image.png

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

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

总结

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

作者:Lin Zhang
来源:GiantPandaCV

推荐阅读

欢迎大家点赞留言,更多Arm技术文章动态请关注极术社区嵌入式AI专栏欢迎添加极术小姐姐微信(id:aijishu20)加入技术交流群,请备注研究方向。

推荐阅读
关注数
18835
内容数
1369
嵌入式端AI,包括AI算法在推理框架Tengine,MNN,NCNN,PaddlePaddle及相关芯片上的实现。欢迎加入微信交流群,微信号:aijishu20(备注:嵌入式)
目录
极术微信服务号
关注极术微信号
实时接收点赞提醒和评论通知
安谋科技学堂公众号
关注安谋科技学堂
实时获取安谋科技及 Arm 教学资源
安谋科技招聘公众号
关注安谋科技招聘
实时获取安谋科技中国职位信息