AI Infra论文阅读之《在LLM训练中减少激活值内存》

最新推荐文章于 2026-01-08 09:30:54 发布

原创

最新推荐文章于 2026-01-08 09:30:54 发布 · 2k 阅读

20 ·

CC 4.0 BY-SA版权

文章标签：

#人工智能 #论文阅读

文章介绍了Megatron-LM中如何通过SequenceParallel和SelectiveActivationRecomputation技术减少大语言模型的激活内存，以及这两种方法如何在1万亿参数模型上实现5倍的记忆减半和90%以上的重计算时间减少。特别是，ContextParallel的引入进一步优化了内存消耗并支持更长序列训练。

写了一个Megatron-LM的3D Parallel进程组可视化的Playground，界面长下面这样：

在这里插入图片描述

可以直接访问：https://huggingface.co/spaces/BBuf/megatron-lm-parallel-group-playground

脚本也开源在：https://github.com/BBuf/megatron-lm-parallel-group-playground 请随意获取和修改。

0x0. 前言

本次阅读一篇Megatron-LM的README贴出的一篇paper，是NVIDIA在2022年上传的，讲的是Megatron-LM里面的Sequence Parallel和Selective Activation Recomputation如何降低大语言模型训练中的激活内存。这里来看一下理论讲解，另外Paper的第4节的激活内存公式估计比较有用。paper链接为：https://arxiv.org/pdf/2205.05198.pdf 。序列并行目前是非常常用的，但是选择性激活重计算可能用得人不多，我想一个很重要的原因应该是 FlashAttention 的出现大大降低了激活值内存大小。但这篇Paper的一些公式仍然是有用处的，至少给我们提供了理论依据。此外，Meagtron-LM推出的Context Parallel是在Sequence Parallel的基础上更近一步，可以继续降低激活值的大小，并且支持更长的序列。

摘要就是说通过Sequence Parallel和Selective Activation Recomputation可以减少激活重计算，把Sequece Parallel和Tensor Parallel结合在一起基本可以避免激活重计算的需要。然后在高达一万亿参数规模的语言模型上评估了上面的两个方法，并展示了这里的放大把激活内存降低了5倍，同时把激活重计算的执行时间开销减少了超过90%。例如，在2240 NVIDIA A100 GPUs上训练一个530B参数的GPT-3风格模型时，实现了54.2%的MFU，比使用重计算实现的42.1%快了29%。

0x1. 介绍

这里简单总结一下，一般来说Meagtron-LM里面张量并行都是放在同一个GPU的节点里面，节点内部由NVLink连接。然后，流水线并行虽然可以减少存储模型参数和优化器状态的内存，但是由于要存储一些Micro Batch的激活，所以并不能减少激活需要的内存。因此，激活内存的存储成为了训练大语言模型的一个关键问题。图1显示了从220亿参数到1万亿参数的四种模型配置所需的内存（模型配置的详细信息在表3中提供）。

在这里插入图片描述

这里的present work就是通过激活重计算（也叫Gradient Checkpointing）来减轻Activation的存储大小。之前标准的做法是在每一个Transformer层的边界进行重计算，paper也把这种方法叫作完全激活重计算。但完全激活重计算会导致增加30-40%左右的时间开销。为了节省这部分计算开销，但又要Scale Up模型，所以就引入了Paper介绍的两种方法Sequence Parallel和Selective Activation Recomputation。

0x2. 相关工作rt

对Megatron-LM的张量并行进行了简单的介绍，没什么干货，忽略。

0x3. Transformer的结构

如下图的图2所示：输入token被送入一个大小为 $v\times h$ 的词嵌入表中，token嵌入与学习到的位置嵌入（大小为 $s\times h$ ）结合，其中 $s$ 是序列长度， $h$ 是隐藏维度， $v$ 是词表大小。嵌入层的输出，即Transformer块的输入，是一个大小为 $s\times b\times h$ 的3维张量，其中 $b$ 是微批量大小。每个Transformer层由一个自注意力块组成，该块有 $a$ 个注意力头，接着是一个增加隐藏大小到 $4 h$ 然后再减少回 $h$ 的两层多层感知器（MLP）。每个Transformer层的输入和输出大小相同，为 $s \times b \times h$ 。最后一个Transformer层的输出被投影回词汇维度以计算交叉熵损失。paper假设词嵌入和输出层权重是共享的。变量名在表1中列出以供参考。

在这里插入图片描述

0x4. Activation Memory

首先，Paper导出了一个近似的公式来估计激活内存的大小，这里的激活指的是在Forward过程中创建并且在Backward中用于梯度计算所必需的任何张量。然后，这里只考虑对激活内存贡献最大的部分不考虑小的Buffer，比如对于LayerNorm来说，输入包含bsh个元素，但均值和方差每个只有sb个元素，由于h一般很大，所以bsh远远大于2sb，所以就忽略掉2sb，以sbh来计算LayerNorm层的激活大小。