本文共 1597 字，大约阅读时间需要 5 分钟。

数据whale干货：大模型显存优化与计算方法

作者：kaiyuan，来源：知乎

摘要

本文详细分析了大模型训练/推理过程中显存的使用情况，探讨了模型参数存储、优化器状态、激活值、梯度值等关键部分对显存消耗的影响，并提出了多种优化方法和计算公式。通过这些方法，可以有效降低大模型训练的显存需求。

---

模型显存内容分析

在大模型的训练或推理过程中，显存的分配主要分为两部分：一部分用于AI框架的运行，另一部分用于系统的底层驱动。从用户侧可以通过显存可视化工具（如PyTorch的显存分析功能）观察训练过程中显存的使用情况。

显存消耗组成

可估算值：模型参数、优化器状态、激活值、梯度值、输出数据。

未命名数据：临时变量、未知数据。

其他：自动梯度数据。

显存消耗差异

由于未知数据的影响，显存估算值与实际测量值可能存在较大差异（误差可超过30%）。例如，估算值为50GB，而实际测试值可能达到75GB。

---

计算公式与优化方法

训练场景

训练过程中，显存消耗主要由以下四部分组成：

模型参数（Model Memory）

模型参数的存储大小与参数量和数据类型有关。例如，1B模型采用fp32存储，约需要3.725GB，LLama13b模型存储空间约为52GB。

优化器状态（Optimizer status）

Adam优化器的状态包括模型副本、Momentum参数和Variance参数。对于8位优化器，存储空间占用减少为4→1→1Bytes。

梯度值（Gradient）

梯度值的存储大小与模型数据类型一致，例如fp32梯度占用为4Bytes。

激活值（Activation）

激活值的存储大小与模型结构和并行策略有关。Megtron论文提供了计算公式：

[ \text{激活值显存} = s \times b \times h \times a \times t \times \lambda ]

其中，s为序列长度，b为微批量大小，h为隐藏层大小，a为注意力头数，t为Tensor并行数值，L为模型层数，λ为精度系数。

并行计算与显存优化

为了降低单卡显存消耗，常用的并行策略包括TensorParallel（TP）、SequenceParallel（SP）、PipelineParallel（PP）和Zero方法。

3D并行

通过将模型、激活值和梯度分割到不同GPU，显存占用可降低至原来的1/3。公式为：

[ \text{优化后的显存} = \frac{\text{单卡显存}}{N} \times (1 - \frac{\text{模型参数}}{\text{总参数}}) ]

重计算（Recomputation）

通过丢弃中间数据，显存占用可降低至原来的1/2。公式为：

[ \text{激活值显存} = s \times b \times h \times a \times \lambda ]

Zero方法

Zero1、Zero2和Zero3通过减少GPU上的参数存储，显存占用可降低至原来的1/8。公式为：

[ \text{Zero方法显存} = N \times \frac{\text{参数总量}}{L} ]

---

推理场景

推理的显存估算公式较为简单，主要包含以下部分： [ \text{总显存} = \text{模型参数} + \text{激活值} + \text{输出数据} ]

---

显存优化总结

显存优化可以通过以下方法实现：

多卡并行：降低单卡显存压力，适合大规模模型训练。

算子优化：选择精度与显存占用平衡的算子。

数据类型优化：采用低精度（如fp16、int8）存储，降低显存占用。

框架优化：消除框架副本，减少内存碎片。

底层API优化：通过优化CUDA函数和显存管理，降低显存消耗。

通过合理搭配以上方法，可以有效降低大模型训练的显存需求，为模型规模的提升提供保障。

---

点赞三连↓

转载地址：http://olwyz.baihongyu.com/