LLM大模型训练/推理的显卡内存需求计算
LLM大模型训练/推理的显卡内存需求计算
在LLM大模型的训练和推理过程中，GPU内存需求是一个关键问题。本文将详细介绍使用单个和多个GPU以及不同的优化器和批处理大小进行LLM训练和推理时GPU内存需求的计算方法。
无论你是从头开始训练LLM、对其进行微调还是部署现有模型，选择合适的GPU对成本和效率都至关重要。在这篇博客中，我们将详细介绍使用单个和多个GPU以及不同的优化器和批处理大小进行LLM训练和推理时GPU内存需求的所有信息。
计算机处理器由多个决定性电路组成，每个电路都可以处于关闭或打开状态。就内存而言，这两种状态由0或1或位表示。一组八位称为一个字节。1个字节可以表示零（00000000）和255（11111111）之间的数字，或28（等于256个不同位置）。通常，在FP-32（包括符号、指数和尾数）数据类型上训练的神经网络最多占用4个字节的内存。
模型参数常用的数据类型如下：
- float（32位浮点）：每个参数4个字节
- half/BF16（16位浮点）：每个参数2个字节
- int8（8位整数）：每个参数1个字节
- int4（4位整数）：每个参数0.5个字节
1、什么会消耗GPU内存？
在模型训练期间，大部分内存被四个东西消耗
1.1 模型参数
模型参数是神经网络的可学习组件。它们定义网络的结构和行为，并在训练期间更新以最小化损失函数。通常，我们有权重和偏差参数。
正如我们已经知道的那样，存储一个数字需要4个字节。假设我们的模型中有P个参数。
- 参数内存（M）= 参数数量（P）x 精度大小（4字节）
- M = Px4
- 16位M = P x 精度大小（2字节）也类似
我们可以添加一个缩放因子并制定一个标准公式，如下所示：
这里1.2表示在GPU内存中加载额外内容的20%开销，Q是加载模型应使用的位数。即16位、8位或4位。
16位Llama 70B需要GPU内存：
这是推理Llama 70b模型所需的总体最低GPU。
1.2 激活
当输入数据通过网络时，激活是每层神经元的中间输出。在前向传递过程中，每层处理输入数据，应用权重、偏差和激活函数（如ReLU、sigmoid等）来产生激活。然后，这些激活将作为下一层输入。
需要存储每个层的激活，因为它们在反向传播期间用于计算梯度。
激活内存 = 激活数量 x 批次大小 x 精度大小
注意：“每个参数的激活”取决于模型架构、层数和序列长度。对于大型模型，激活通常需要与参数相当或超过参数的内存。将序列长度加倍也可能使激活内存加倍。
近似值：没有固定的公式来计算激活的GPU内存。对于较大的模型，激活所需的内存可能大致与参数的内存相似或略大。
1.3 梯度
梯度是损失函数关于模型参数的偏导数。它们表示应调整每个参数多少以最小化损失函数。
在反向传播期间，损失通过网络向后传播，并计算每个参数（权重和偏差）的梯度。优化器使用这些梯度来更新参数，从而减少整体损失。
存储梯度所需的内存等于参数本身所需的内存。由于每个参数都有相应的梯度，因此它们的内存要求相同。
梯度内存 = 参数内存
1.4 优化器状态
优化器状态是某些优化算法（如Adam、RMSprop）维护的附加变量，用于提高训练效率。这些状态有助于根据过去的梯度更新模型参数。
不同的优化器维护不同类型的状态。例如：
- SGD（随机梯度下降）：没有附加状态；仅使用梯度来更新参数。
- Adam：为每个参数维护两个状态：一阶矩（梯度平均值）和二阶矩（梯度平方平均值）。这有助于动态调整每个参数的学习率。对于具有100万个参数的模型，Adam需要为每个参数维护2个附加值（一阶矩和二阶矩），从而产生200万个附加状态。
优化器状态的内存 = 参数数量 x 精度大小 x 优化器乘数
2、单GPU内存需求计算
我们举个例子
我们想在单个GPU上以混合精度（2字节）训练100亿模型。
- 参数内存=参数数量 x 2字节 (FP16)
- 参数内存=10B x 2字节 = 20 GB
- 激活内存=每个参数的激活 x 批次大小 x 精度大小
我们可以计算每层激活内存，而不是计算激活的总内存，这是一种高效的方法，需要的内存更少，因为它可以在下一层使用。
- 每层神经元的近似数量 = sqrt(10B) ≈ 每层100k个神经元
- 一层的激活内存 ≈ 32 x 100k x 2字节 ≈ 每层6.4 MB
对于大型模型中的层（假设有数百层），激活内存最多可达数十GB。
因此，正如我们之前讨论过的，对于32的批次大小，大约需要20-40GB的内存。现在，如果我们将批次大小加倍，这个范围可以加倍。
- 梯度内存 = 参数内存
- 梯度内存 = 20 GB
- 优化器状态内存 = 参数数量 x 4字节 (FP32) x 2 (用于Adam)
- 优化器状态内存 = 10B x 4字节 x 2 = 80 GB
总内存估计：
- 参数内存：20 GB
- 激活内存：≈20–40 GB（取决于批次大小）
- 梯度内存：20 GB
- 优化器状态内存：80 GB
- 总内存 = 20 + 20 到 40 + 20 + 80 = 140 到 160 GB
3、多个GPU的内存计算
要计算在n个GPU上训练时每个GPU的内存需求，我们需要考虑如何使用数据并行和模型并行等并行技术在GPU上分配内存。
关键假设：
- 模型并行：模型的参数在GPU之间分配，因此每个GPU仅存储总模型参数的一小部分。梯度和优化器状态也同样被划分。
- 数据并行：每个GPU都会获得整个模型参数的副本，但数据批次会在GPU之间分配。激活是针对每个GPU的小批次单独计算的。
如果我们使用模型并行性，那么所有模型参数、梯度和优化器统计数据都是分布式的。但是，每个GPU仍然需要存储其批次部分的激活。激活的内存不会随着GPU数量的增加而减少，因为每个GPU都独立处理自己的数据。
因此，对于所有GPU来说，激活所需的内存仍然相同
因此，在n个GPU上以混合精度（2字节）训练100亿模型所需的总内存为：
如果我们想使用2个GPU训练LLM，我们需要大约80到100GB的内存。
原文链接：LLM显卡内存需求计算 - BimAnt