[LLM]大模型显存计算公式与优化

[LLM]大模型显存计算公式与优化

AI算法在服务器中运行时，一个常见问题“单张GPU能承载多少模型参数？”，该问题跟模型结构、引擎框架、驱动版本、GPU硬件相关。本文围绕大模型的训练/推理场景，介绍transformer类模型的显存计算公式，帮助读者能更好的了解全局显存的组成以及如何优化显存。文中涉及的主要问题：

• 如何有效估算一个模型加载后的显存值？
• 计算值与实际GPU中的最大值的差距可以有多大？
• 大模型切分策略是如何降低显存的？计算公式怎么构建？
• 优化显存的方法和常见的优化思路？

01**** 模型显存内容分析

在模型训练/推理时，显存（显卡的全局内存）分配一部分是给AI框架，另一部分给了系统（底层驱动）。总的显存消耗量可以通过API查询，比如在NVIDIA-GPU上通过nvidia-smi指令能够打印出各个进程的显存消耗量。

+---------------------------------------------------------------------------------------+
| Processes:                                                                            |
|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |
|        ID   ID                                                             Usage      |
|=======================================================================================|
|    1   N/A  N/A     67321      C         .../anaconda3/envs/py/bin/python    23646MiB |
|    1   N/A  N/A     71612      C         .../anaconda3/envs/py/bin/python      848MiB |
|    2   N/A  N/A     67321      C         .../anaconda3/envs/py/bin/python    25776MiB |
+---------------------------------------------------------------------------------------+

其中系统层的显存消耗一般由驱动控制，用户不可控；框架侧的显存消耗用户可控，也是本文分析的重点。以PyTorch框架为例通过显存可视化工具，看一下训练过程中显存的消耗。如下图是一个模型训练过程中已用显存的数值随时间的变化

注意：数据是具体的消耗值，不等于cudaMalloc创建的显存值。

显存消耗的内容包括：模型参数（parameter）、优化器状态值（optimizer_state）、激活值（activation）、梯度值（gradient）、输出数据（input）、临时变量（temporary）、自动梯度（autograd_detail）、未知数量（unknown)。从用户侧可以将这些数据进行一个分类：

• 可估算值：模型参数（parameter）、优化器状态值（optimizer_state）、激活值（activation）、梯度值（gradient）、输出数据（input）
• 未命名数据：临时变量（temporary）、未知数据（unknown)
• 其他（框架）：自动梯度（autograd_detail）

其中“未命名数据”来源可能是用户创建的一些临时变量，这些变量未参与图的计算过程，所以未被统计；或者是一些未被框架跟踪（tracing）到的数据。“自动梯度数据"是在反向传播求解梯度时产生的一些变量；

我们在显存计算时会发现“为什么有时显存估算值和实际测量值相差较大？”其中一个可能的原因是：未知的数据太大。即显存中可估算值占比相对较小，其它不可估算值的数据占比较大，导致计算值和实际值差距较大（误差可超过30%），比如估算得到的显存消耗为50GB，而实际测试达到了75GB。

如下图是运行一个LLM模型采集的一些过程数据，可以看到unknown占比有时能达到30%。

不同时刻显存的占比变化

0****2 计算公式

2.1 训练场景

训练显存消耗（可估算部分）主要包括：模型参数（Model）+ 优化器状态（Optimizer status）+梯度值（Gradient）+激活值（Activation）。根据数值的变化，可将显存消耗分为静态/动态值。训练过程中，模型参数、优化器状态一般不会变化，这两部分归属于静态值；激活值、梯度值会随着计算过程发生变化，将它们归类到动态值。

下面主要来看一下这四种类型值的估算方法：

2.1.1 模型显存（Model Memory)

模型自身所占用的显存大小与参数量、参数类型相关。常见类型fp32、fp16/bf16、还有int8、fp8等。

计算满足

根据不同数据类型，有如下计算公式（ 单位：GB）：

关于模型保存的大小估算方法：

存储checkpoint（ckpt）时仅考虑模型本身，只要将显存上模型内容存储到磁盘中。举例：以1B（billion）模型为例，若采用fp32类型将其存储在磁盘上，其大小为：

1B模型需要3.725GB存储空间，进一步近似认为1B≈****4GB，可方便作存储的估算推导，如LLama13b，大约需要52GB存储空间。注意：混合精度（Mixed-precision）最后存储的类型也是fp32，公式也适合混合精度。

2.1.2 优化器状态（Optimizer status）

在LLM中常见的优化器是Adam，优化器中每个参数需要一个Momentum和一个Variance状态参数，在混合精度训练中Adam还有一份模型参数副本。Adam参数器状态值计算公式（单位GB）：

其中（4+4+4）的内容：

• 模型副本 4 Bytes
• Momentum 参数 4 Bytes
• Variance 参数 4 Bytes

如果是8位优化器，则计算变为：

• 模型副本 4 Bytes
• Momentum参数 1Byte
• Variance参数 1Byte

2.1.3 梯度值（Gradient）

梯度值与模型数据类型保持一致，计算如下（单位GB）：

1.fp32的模型梯度值： 

2.fp16或者bf16模型：

2.1.4 激活值（Activation）

激活值的大小跟模型参数、重计算、并行策略等相关，这里我们参考Megtron论文里面给的计算公式，来求解激活值所占用的显存大小。

激活值显存消耗：

(单位GB)，参数说明：

• s   序列长度（sequence length), tokens的量
• b   微批量大小（microbatch size）
• h   隐藏层大小（hidden dimension size）
• a   attention的头数 （number of attention heads）
• L    transformer模型的层数
• γ   比例系数，当为fp16时 值等于1 / (1024 * 1024 * 1024)。

2.2 训练的并行计算公式 目前，单卡的物理显存基本不能满足大模型的训练需求，一般会采用模型并行方式来降低单卡显存消耗，常见的几种方法：TP/SP/PP/Zero/重计算，这些方法出现在DeepSpeed、Megtron等并行框架中，目标都是让GPU能够装下更大的模型，其中：

• TP（TensorParallel）：tensor并行；
• SP（SequenceParallel）：序列并行；
• PP（PipelineParallel）：pipeline并行；
• Zero： 参数服务器，分为Zero1/2/3，最早出现在deepspeed中

当没有并行策略时，仅模型本身的显存需求（单卡）计算如下：

经过并行策略的调整，显存需求可变为（举例，PP/TP/zero1）：

2.2.1 3D并行

3D并行主要是TP(SP)/PP/DP，其中DP为数据并行主要用于提升bs（batch size），DP不降低单卡的显存消耗，但TP(SP)/PP/DP存在一个耦合关系，DP的设置一般满足：DP=GPUs/(TP*PP） 。而TP(SP)/PP可降低模型、激活值、梯度的显存占用大小。

3D并行对显存计算的影响计算：

• 模型显存：

• 梯度显存：

注意：梯度显存没有除以TP，主要是考虑到反向计算时需要AllGather出完整gradient。

3D对激活值显存的消耗改变需要结合重计算公式进一步分析。另一个问题，当前比较流行的MoE方式也会改变模型的参数分布进而改变计算。但认为MoE构造的是多个小模型，改变的是模型的结构，这里计算暂不展开。

考虑MoE时参数的变化

2.2.2 重计算（Recomputation）

一般而言，我们会把前向计算中的中间数据保存下来用于反向计算，从而避免反复计算。而重计算是指为了降低显存消耗先丢弃一些前向计算结果，在反向传播时再重新计算得到。

结合论文[Reducing Activation Recomputation in Large Transformer Models]里面给的计算公式，激活值所占用的显存的计算公式如下：

1.无重计算无并行： 

2.Tensor并行（基线）：

3.Tensor + 序列并行： 

4.选择重计算、tensor并行：

5.选择重计算、tensor + 序列并行：

6.全部重计算:

单位GB，参数说明：

• s : 序列长度（sequence length), tokens的量
• b : 微批量大小（microbatch size）
• h : 隐藏层大小（hidden dimension size）
• a : attention的头数 （number of attention heads）
• t : tensor并行数值（tensor parallel size）
• L : transformer模型的层数
• γ : 比例系数，当为fp16时 值等于1 / (1024 * 1024 * 1024)。

假设我们选用Tensor和序列并行、不开重计算，则单卡的公式变为：

其中，

2.2.3 Zero方法

Zero方法对显存的优化和原理参考其论文[arxiv.org/abs/1910.02…

zero策略下显存消耗的计算变化

假设不考虑3D并行和重计算，开启Zero的计算公式为：

其中N 是GPU的数量；LiveParams是Zero3引入的参数，这些参数用于控制模型中哪些参数需要加载在GPU中，本身的显存占用不可忽视。

2.2.4 训练的综合计算列举

当条件确定好后，我们可将上述的公式综合起来求解总的显存消耗。通过一个具体的示例来说明，假设相关的运算条件：

1. 采用混合精度训练；
2. 开启TP/SP/PP；
3. 不开重计算；
4. 开启Zero2；

混合精度的单层的数据配置一般如下图所示，需要注意的是master weights只要算一次，要么在优化器中计算要么在模型中计算，这里默认在优化器中考虑。

混合精度数值类型

计算公式如下（单位GB）：

（

其中：

相关参数说明：

• params: 模型参数
• N：GPU数量
• PP：pipeline并行数值
• TP : Tensor + 序列并行 数值
• s : 序列长度（sequence length), tokens的量
• b : 微批量大小（microbatch size）
• h : 隐藏层大小（hidden dimension size）
• a : attention的头数 （number of attention heads）
• t : tensor并行数值（tensor parallel size）
• L : transformer模型的层数

注意：公式计算得到是一个估算值，且只考虑了模型部分，实际运行中的总数还需要考虑框架、分布式通信库、环境变量、算法产生副本数据。

2.3 推理场景

推理的显存量组成成分比训练简单，有一个简单的估算公式：

总显存占用：

相关内容可参看这篇blog：Transformer Inference Arithmetic | kipply's blog(kipp.ly/transformer…)

总之，通过综合求解公式可以知道模型显存消耗主要部分，能帮助我们确定显存的优化的策略。

0****3 显存优化

由于大模型的参数成倍数的增长，远超出了单GPU物理显存所能承载的范围，大模型训练必然需要进行显存优化。显存优化要么是优化算法本身，降低模型算法的显存消耗；要么是去扩大显存，通过一些置换方式获得“额外“空间，由于显存物理大小一定，我们获得额外空间的方式不外乎两种：

• **时间换空间；**如，重计算
• **空间转移；**如，多卡并行/offload

其中，时间换空间通常会消耗算力、带宽；空间转移主要是消耗I/O带宽，有一定的时延，可能会降低吞吐**。**

显存优化的过程一般是从模型算法本身到底层，可以参考的优化路径：

多卡并行  ->  算子/数据类型 -> 消除框架副本 ->  显存管理  -> 底层API

1、多卡并行：该手段相对来说是使用频率最高，且一般不会影响运算的精度，可以用2节中的计算公式为参考去设计新的TP/PP/DP/Zero/重计算的相关参数来降低显存消耗。缺点：这些方式可能会增加额外的带宽消耗。

2、算子优化：选取精度相同但显存消耗更低的算子/方案。缺点：一般情况下，算子优化的过程耗时较长。

3、数据类型修改：用低精度替换高精度数据。比如用fp16代替fp32，或者用更低的int8/int4。缺点：该方式可能影响训练收敛性/推理性能。

4、消除框架副本：在AI框架（如pytorch）中有些数据是一些由框架产生的中间副本，可以进行优化消除；缺点：游湖成本较大。

5、显存管理：通过显存管理的知识可知[PyTorch显存管理]，框架的显存管理会产生显存碎片，通过优化显存管理来优化碎片；缺点：目前可用的手段较少。

6、底层API： 在GPU的驱动库中/CUDA算子库中，不同API显存消耗不一样，我们可以用显存消耗更小算子去替换大显存消耗算子，比如FlashAttention；有些默认的操作会产生额外系统显存，也可以考虑替换更高版本优化后的API。

---