大模型训练需要更大的算力

大模型训练所需的总算力，其实很简单，6 * 模型的参数量 * 训练数据的 token 数就是所有训练数据过一遍所需的算力。这里的 6 就是每个 token 在模型正向传播和反向传播的时候所需的乘法、加法计算次数。

一堆矩阵相乘，简单来想就是左边若干个神经元，右边若干个神经元，组成一个完全二分图。选出其中任意一个左边的神经元 L 和右边的神经元 R。

1. 正向传播的时候： L把它的输出乘上 L和 R 之间的权重 w，发给 R；R不可能只连一个神经元吧，总要把多个 L的加到一起，这就是 reduce，需要一次加法。
2. R把它收到的梯度乘上 L和 R 之间的权重 w，发给 L；L也不可能只连一个 R，需要把梯度 reduce 一下，做个加法；别忘了权重 w 需要更新，那就要计算 w 的梯度，把 R 收到的梯度乘上 L正向传播的输出（activation）；一个 batch 一般有多个 sample，权重 w 的更新需要把这些 sample 的梯度加到一起。
3. 一共 3 次乘法，3 次加法，不管 Transformer 多复杂，矩阵计算就是这么简单，其他的向量计算、softmax 之类的都不是占算力的主要因素，估算的时候可以忽略。

有了模型训练所需的总算力，除以每个 GPU 的理论算力，再除以 GPU 的有效算力利用比例，就得到了所需的 GPU-hours。

大模型训练需要更多的显存内存资源

深度学习训练需要的内存包括模型参数、反向传播的梯度、优化器所用的内存、正向传播的中间状态（activation），显存占用 = 模型参数大小 + Batch Size * 优化器参数与中间变量的大小。

1. 优化器所用的内存的计算其实也很简单，如果用最经典的 Adam 优化器，它需要用 32 位浮点来计算。即使我们使用mixed-precision进行计算，每个参数需要也要存 4 字节的 32 位版本（正向传播时用 16 位版本，优化时用 32 位版本），还需要存 4 字节的 momentum 和 4 字节的 variance，一共 12 字节。如果是用类似 SGD 的优化器，可以不存 variance，只需要 8 字节。
2. 正向传播的中间状态（activation）是反向传播时计算梯度必需的，而且跟 batch size 成正比。Batch size 越大，每次读取模型参数内存能做的计算就越多，这样对 GPU 内存带宽的压力就越小。划重点**：正向传播的中间状态数量是跟 batch size 成正比的**。

大模型推理的内存计算只需考虑模型权重即可。

大模型训练的内存计算往往会考虑包括模型权重、反向传播的梯度、优化器所用的内存、正向传播的激活状态内存。

接下来以ChatGLM-6B为例，它的参数设置为隐藏层神经元数量（hidden_size）为 4096，层数（num_layers）为 28，token 长度为 2048，注意力头数(attention heads）为 32，讲解怎么计算推理内存和训练内存。

推理内存

模型权重

• 对 int8 而言，模型内存 =1 * 参数量 （字节）
• 对 fp16 和 bf16 而言，模型内存=2 * 参数量（字节）
• 对 fp32 而言，模型内存= 4 * 参数量（字节）

因为 1 GB ≈ 1B字节，也正好和1B参数量的数据量级一致，估算时就比较简单了。所以对于一个ChatGLM-6B而言，就是：

• 对 int8 而言，模型内存=1 * 6GB=6GB
• 对 fp16 和 bf16 而言，模型内存=2 * 6GB=12GB
• 对 fp32 而言，模型内存=4 * 6GB=24GB

推理总内存

除了用于存储模型权重的内存外，在实际的前向传播过程中还会产生一些额外的开销。根据经验，这些额外开销通常控制在总内存的20%以内。

因此，推理总内存≈1.2×模型内存

训练

模型权重

可以使用纯 fp32 或纯 fp16 训练模型:

• 纯 fp32，模型内存=4 * 参数量（字节）
• 纯 fp16，模型内存=2 * 参数量（字节）

除了常规推理中讨论的模型权重数据类型，训练阶段还涉及混合精度训练。

混合精度 (fp16/bf16 + fp32), 模型内存=2 * 参数量（字节）

对于一个ChatGLM-6B而言，就是：

• 纯 fp32，模型内存=4 * 6GB=24GB
• 纯 fp16，模型内存=2 * 6GB=12GB

混合精度 (fp16/bf16 + fp32), 模型内存=2 * 6GB=12GB

优化器状态

对于纯 AdamW，优化器内存=12 * 参数量（字节）

对于像 bitsandbytes 这样的 8 位优化器，优化器内存=6 * 参数量（字节）

对于含动量的类 SGD 优化器，优化器内存=8 * 参数量（字节）

• 对于一个ChatGLM-6B而言，就是：
• 对于纯 AdamW，优化器内存=12 * 6GB=72GB
• 对于像 bitsandbytes 这样的 8 位优化器，优化器内存-6 * 6GB=36GB
• 对于含动量的类 SGD 优化器，优化器内存=8 * 48GB=36GB

梯度

梯度可以存储为 fp32 或 fp16 (梯度数据类型通常与模型数据类型匹配。因此在 fp16 混合精度训练中，梯度数据类型为 fp16)

• 对于 fp32，梯度内存=4 * 参数量（字节）
• 对于 fp16，梯度内存=2 * 参数量（字节）
• 对于一个ChatGLM-6B而言，就是：
• 对于 fp32，梯度内存=4 * 6GB=24GB
• 对于 fp16，梯度内存=2 * 6GB=12GB

激活状态

在进行LLM（大语言模型）训练时，现代GPU通常会遇到内存问题，而不是算力问题。因此，激活重计算（也称为激活检查点）变得非常流行，它是一种以计算力为代价来减少内存使用的方法。激活重计算/检查点的主要思路是重新计算某些层的激活，而不将它们存储在GPU内存中，从而降低内存使用量。具体来说，减少内存的多少取决于我们选择重新计算哪些层的激活。 https://blog.eleuther.ai/transformer-math/ 接下来，假设激活数据类型为 fp16，没有使用序列并行

• 无重计算的激活内存=token 长度 * batch size * hidden layer 的神经元数量 * 层数(10+24/t+5 * a * token 长度/hidden layer 的神经元数 * t) 字节
• 选择性重计算的激活内存=token 长度 * batch size * hidden layer 的神经元数量 * 层数(10+24/t) 字节
• 全部重计算的激活内存=2 * token 长度 * batch size * hidden layer 的神经元数量 * 层数 字节

其中：

• a 是 transformer 模型中注意力头 (attention heads) 的个数
• t 是张量并行度 (如果无张量并行，则为 1)

对于一个ChatGLM-6B而言，就是：

token 长度 * batch size * hidden layer 的神经元数量 * 层数 = 2048 * 1 * 4096 * 28 ≈ 0.23GB

无重计算的激活内存 = 0.23GB * (10+24/1+5 * 32 * 2048/4096 * 1) = 0.23 * 114 = 26.22G

选择性重计算的激活内存 = 0.23GB * (10+24/1) = 7.8G

全部重计算的激活内存 = 2 * 0.23GB = 0.46GB

由于重计算的引入也会引起计算成本的增加，具体增加多少取决于选择了多少层进行重计算，但其上界为所有层都额外多了一次前向传播，因此，更新后的前向传播计算成本如下: https://blog.eleuther.ai/transformer-math/ 2 * token数 * 模型参数 ≤ C（前向传播）≤ 4 * token数 * 模型参数

总结

因为训练大模型时通常会采用AdamW优化器，并用混合精度训练来加速训练，所以训练一个ChatGLM-6B所需的训练总内存为：

训练总内存=模型内存+优化器内存+激活内存+梯度内存 = 12GB + 72GB + 12Gb + 7.8GB = 103GB

训练一个大模型需要多少显存？ (opens new window)

1、一张1GB显存单位转换

• 1GB=1024MB=1024*1024KB=1024*1024*1024Byte（字节）=1024*1024*1024*8Bit

参数对应表：

参数类型	Byte（字节数）
FP32	4
FP16、BF16	2
INT8	1
INT4	0.5

2、显存会被以下四个方面占用

2.1、模型参数W：模型的权重，偏置，位置编码和attn_mask参数等。所有这些参数都需要存储在显存中。Qwen2-VL-2B-Instruct，参数量2B，以BF16方式存储，模型参数大概需要占用4GB的显存空间，计算方式：

$W=2B=21000^ 3BF16=21000^3*2Byte\approx4GB $ ,

2.2、参数梯度（全参数训练中）

以BF16方式训练，参数梯度需要4G，计算方式：

$G=2B(BF16)=2\ast {1000}^{3}BF16=2\ast {1000}^3\ast 2Byte\approx 4GB$

2.3、中间参数（激活值，随模型输入大小，dim，代码实现方式发生变化，若batch_size=1，一个样本1000个token大概需要模型参数的0.7倍-1倍）：

• 输入数据：input=input.to("cuda")。
• 标签数据：label=label.to("cuda")。
• 注意力分数：attn = (q @ k.transpose(-2, -1))。
• 保留残差数据：resnet_shortcut = x，等中间参数。

2.4、优化器（模型参数副本W和优化器状态值，可能存储梯度G，默认以FB32存储）在反向传播过程中，需要模型参数副、求解梯度的相关值或梯度G。例如优化器Adam，求解梯度会保存梯度动量项V、平方梯度S的平均值等。以Qwen2-VL-2B-Instruct，参数量2B使用Adam梯度优化算法训练为例，保存V和S参数，需要以FB32存储。在训练过程中分别需要8GB，共计16GB，计算方式如下：

$V=S=2B(BF32)=21000^ 3BF32=21000^3*4Byte\approx8GB $

模型参数副需要以FB32存储，在训练过程中需要8GB，计算方式如下：

$W^A=2B(BF16)=2\ast {1000}^{3}BF32=2\ast {1000}^3\ast 4Byte\approx 8GB$

如果使用deepspeed-zerox方式训练，如果有N张显卡，梯度优化器上的参数会被才分成N份，分别保存在不同显卡上，如果一张显卡公式如下：

$W^A=G^A=V=S=2B(BF32)=（2\ast {1000}^{3}BF32）/N=2\ast {1000}^3\ast 4Byte/N\approx 8GB/N$

Adam公式

2.5 训练过程中合计

普通模式： $W+G+W^A+V+S+0.7W=24.8G$

以deepspeed方式训练，8卡： $W+G+(W^A+G^A+V+S)/8+0.7W=14.8G$

注意**：SGD不需要保存V和S，SGDM不需要保存S，AdaGrad不需要保存V。**导航栏

【一起学习深度学习】 深度学习导航 | 快速复习笔记 | 知识点总结64 赞同 · 34 评论文章64 赞同 · 34 评论文章 (opens new window)如果有疑问，建议，错误，欢迎留言大家指正。。。

更新不易啊，欢迎点赞关注。。。

内容还在持续更新中。。。