本文次要介绍用于估算 transformer 类模型计算量需要和内存需要的相干数学方法。
引言
其实，很多无关 transformer 语言模型的一些根本且重要的信息都能够用很简略的办法估算进去。可怜的是，这些公式在 NLP 社区中鲜为人知。本文的目标是总结这些公式，说明它们是如何推导进去的及其作用
留神: 本文次要关注训练老本，该老本次要由 GPU 的 VRAM 主导。如果你想晓得无关推理老本 (通常由推理提早主导) 的信息，能够读读 Kipply 写的 这篇精彩博文。
算力要求
下式可用于估算训练一个 transformer 模型所需的算力老本:
$$ C \approx \tau T = 6PD $$
这里:
- $C$ 是训练 transformer 模型所需的计算量，单位为总浮点运算数 (FLOP)
- $C=C_{\text {前向}}+C_{\text {后向}}$
- $C_{\text {前向}} \approx 2PD$
- $C_{\text {后向}} \approx 4PD$
- $\tau$ 是训练集群的理论总吞吐量: $\tau=\text {GPU 数} \times \text {每 GPU 的理论每秒浮点运算数 (理论 FLOPs) }$，单位为 FLOPs
- $T$ 是训练模型所破费的工夫，以秒为单位
- $P$ 是 transformer 模型的参数量
- $D$ 是数据集大小，示意为数据集的总词元数
该式由 OpenAI 的缩放定律论文 和 DeepMind 的缩放定律论文 提出并经其试验验证。想要获取更多信息，可参阅这两篇论文。
上面，咱们讨论一下 $C$ 的单位。$C$ 是总计算量的度量，咱们能够用许多单位来度量它，例如:
FLOP - 秒，单位为 ${\text {每秒浮点运算数}} \times \text {秒}$GPU - 时，单位为 $\text {GPU 数}\times\text {小时}$- 缩放定律论文偏向于以
PetaFLOP - 天为单位，即单位为 $10^{15} \times 24 \times 3600$
这里须要留神 $\text {理论 FLOPs}$ 的概念。尽管 GPU 的规格书上通常只宣传其实践 FLOPs，但在实践中咱们从未达到过这些理论值 (尤其在分布式训练时！)。咱们将在计算成本这一大节列出分布式训练中常见的 $\text {理论 FLOPs}$ 值。
请留神，下面的算力老本公式来自于 这篇对于 LLM 训练老本的精彩博文。
参数量与数据集的衡量
严格来讲，你能够得心应手地应用任意数量的词元来训练 transformer 模型，但因为参加训练的词元数会极大地影响计算成本和最终模型的性能，因而须要小心衡量。
咱们从最要害的局部开始谈起:“计算最优”语言模型。“Chinchilla 缩放定律”，得名于提出“计算最优”语言模型论文中所训练的模型名，指出计算最优语言模型的 参数量 和 数据集大小 的近似关系满足: $D=20P$。该关系成立基于一个前提条件: 应用 1,000 个 GPU 1 小时和应用 1 个 GPU 1,000 小时老本雷同。如果你的状况满足该条件，你应该应用上述公式去训练一个性能最优且 GPU - 时 老本最小的模型。
但 咱们不倡议在少于 200B 词元的数据集上训练 LLM。 尽管对于许多模型尺寸来说这是“Chinchilla 最优”的，但生成的模型通常比拟差。对于简直所有利用而言，咱们倡议确定你可承受的推理老本，并训练一个满足该推理老本要求的最大模型。
计算成本的经验值
Transformer 模型的计算成本通常以 GPU - 时 或 FLOP - 秒 为单位。
- GPT-NeoX 的
理论 TFLOP/s在失常注意力机制下达到 150 TFLOP/s/A100，在 Flash 注意力机制下达到 180 FLOP/s/A100。这与其余高度优化的大规模计算库统一，例如 Megatron-DS 的值是在 137 和 163 TFLOP/s/A100 之间。 - 一个通用的教训法令是
理论 TFLOP/s可至 120 TFLOP/s/A100 左右。如果你失去低于 115 TFLOP/s/A100 的值，可能是你的模型或硬件配置有问题。 - 借助 InfiniBand 等高速互连设施，你能够在数据并行维度上实现线性或亚线性扩大 (即减少数据并行度应该以近乎线性的形式减少整体吞吐量)。下图显示了在橡树岭国家实验室 (Oak Ridge National Lab) 的 Summit 超级计算机上测试出的 GPT-NeoX 库的扩展性。请留神，这张图用的是 V100，而本文中的其余大多数例子都是基于 A100 的。
内存需要
Transformer 模型通常由其 参数尺寸 来形容。然而，依据给定一组计算资源确定要训练哪些模型时，你须要晓得 该模型将占用多少空间 (以字节为单位)。这不仅须要思考你的本地 GPU 能够推理多大的模型，还须要思考给定训练集群中的总可用 GPU 内存可供训练多大的模型。
推理
模型权重
大多数 transformer 模型都应用 混合精度 进行训练，能够是 fp16 + fp32 或是 bf16 + fp32。混合精度升高了模型训练和推理所需的内存量。推理时，咱们还能够将语言模型从 fp32 转换为 fp16 甚至 int8，而没有实质性的精度损失。上面咱们看看在不同的数据类型下，模型所需内存有什么不同 (以字节为单位):
- 对 int8 而言，$\text {模型内存}=1 \text {字节} /\text {参数}\cdot \text {参数量}$
- 对 fp16 和 bf16 而言，$\text {模型内存}=2 \text {字节} /\text {参数} \cdot \text {参数量}$
- 对 fp32 而言，$\text {模型内存}=4 \text {字节} /\text {参数}\cdot \text {参数量}$
推理总内存
除了存储模型权重所需的内存外，理论中前向流传过程中还会有大量额定开销。依据咱们的教训，此开销在 20% 以内，该比例通常与模型无关。
总的来说，答复“这个模型是否适宜推理”这一问题，能够用下述公式来取得不错的预计:
$\text {推理总内存} \approx 1.2 \times \text {模型内存}$
本文不会深究该开销的起源，留待前面的文章来论述。在本文的后续局部，咱们将次要关注模型训练的内存。如果你有趣味理解更多无关推理计算需要的信息，请查看 这篇深刻介绍推理的精彩博文。当初，咱们要开始训练了！
训练
除了模型权重之外，训练还须要在设施内存中存储优化器状态和梯度。这就是为什么当你问“我须要多少内存来运行模型？”，他人会立刻答复“这取决于是训练还是推理”。训练总是比推理须要更多的内存，通常多得多！
模型权重
首先，能够应用纯 fp32 或纯 fp16 训练模型:
- 纯 fp32，$\text {模型内存}=4 \text {字节} /\text {参数} \cdot \text {参数量}$
- 纯 fp16，$\text {模型内存}=2 \text {字节} /\text {参数} \cdot \text {参数量}$
除了推理中探讨的常见模型权重数据类型外，训练还引入了 混合精度 训练，例如 AMP。该技术寻求在放弃收敛性的同时最大化 GPU 张量核的吞吐量。古代 DL 训练畛域常常应用混合精度训练，因为: 1) fp32 训练稳固，但内存开销高且不能利用到 NVIDIA GPU 张量核、2) fp16 训练稳固但难以收敛。更多混合精度训练相干的内容，咱们倡议浏览 tunib-ai 的 notebook。请留神，混合精度要求模型的 fp16/bf16 和 fp32 版本都存储在内存中，而模型须要的内存如下:
- 混合精度 (fp16/bf16 + fp32), $\text {模型内存}=2 \text {字节} /\text {参数} \cdot \text {参数量}$
正如下面所讲，这个仅仅是模型内存，还须要额定加上 $4 \text {字节 / 参数} \cdot \text {参数量}$ 的 用于优化器状态计算 的模型正本，咱们会在上面的 优化器状态 一节中算进去。
优化器状态
Adam 有奇效，但内存效率非常低。除了要求你有模型权重和梯度外，你还须要额定保留三个梯度参数。因而，
-
对于纯 AdamW，$\text {优化器内存}=12 \text {字节}/\text {参数}\cdot \text {参数量}$
- fp32 主权重: 4 字节 / 参数
- 动量 (momentum): 4 字节 / 参数
- 方差 (variance): 4 字节 / 参数
-
对于像 bitsandbytes 这样的 8 位优化器，$\text {优化器内存} =6 \text {字节} /\text {参数} \cdot \text {参数量}$
- fp32 主权重: 4 字节 / 参数
- 动量: 1 字节 / 参数
- 方差: 1 字节 / 参数
-
对于含动量的类 SGD 优化器，$\text {优化器内存} =8 \text {字节} /\text {参数} \cdot \text {参数量}$
- fp32 主权重: 4 字节 / 参数
- 动量: 4 字节 / 参数
梯度
梯度能够存储为 fp32 或 fp16 (请留神，梯度数据类型通常与模型数据类型匹配。因而，咱们能够看到在 fp16 混合精度训练中，梯度数据类型为 fp16)，因而它们对内存开销的奉献为:
- 对于 fp32，$\text {梯度内存} = 4 \text {字节} /\text {参数} \cdot \text {参数量}$
- 对于 fp16，$\text {梯度内存} = 2 \text {字节} /\text {参数} \cdot \text {参数量}$
激活和 Batch Size
对于 LLM 训练而言，古代 GPU 通常受限于内存瓶颈，而不是算力。因而，激活重计算 (activation recomputation，或称为激活检查点 (activation checkpointing) ) 就成为一种十分风行的以计算换内存的办法。激活重计算 / 检查点次要的做法是从新计算某些层的激活而不把它们存在 GPU 内存中，从而缩小内存的使用量。内存的缩小量取决于咱们抉择革除哪些层的激活。举个例子，Megatron 的选择性重计算计划的成果如下图所示:
图中，红色虚线示意 A100-80GB GPU 的内存容量，“present work”示意应用选择性激活重计算后的内存需要。请参阅 Reducing Activation Recomputation in Large Transformer Models 一文，理解更多详细信息以及下述公式的推导过程。
上面给出存储 transformer 模型激活所需内存的根本公式:
$$
\begin {align*}\text {无重计算的激活内存}=sbhL (10+\frac {24}{t}+5\frac {a \cdot s}{h\cdot t}) \text {字节}\end {align*}
$$
$$
\begin {align*}\text {选择性重计算的激活内存}=sbhL (10+\frac {24}{t}) \text {字节}\end {align*}
$$
$$
\begin {align*}\text {全部重计算的激活内存}=2 \cdot sbhL \text {字节}\end {align*}
$$
其中:
- $s$ 是序列长度，即序列中词元的个数
- $b$ 是每个 GPU 的 batch size
- $h$ 是每个 transformer 层的隐含维度
- $L$ 是 transformer 模型的层数
- $a$ 是 transformer 模型中注意力头 (attention heads) 的个数
- $t$ 是张量并行度 (如果无张量并行，则为 1)
- 咱们假如没有应用序列并行
- 咱们假如激活数据类型为 fp16
因为重计算的引入也会引起计算成本的减少，具体减少多少取决于抉择了多少层进行重计算，但其上界为所有层都额定多了一次前向流传，因而，更新后的前向流传计算成本如下:
$2PD\leq C\_{\text {forward}}\leq4PD$
训练总内存
至此，咱们失去了答复“这个模型是否适宜训练”这一问题的一个很好的估算公式:
$$
\begin {align*}\text {训练总内存} = \text {模型内存}+ \text {优化器内存}+ \text {激活内存}+ \text {梯度内存}\end {align*}
$$
分布式训练
分片优化器 (sharded optimizer)
微小的优化器内存开销促使大家设计和实现了分片优化器，目前罕用的分片优化器实现有 ZeRO 和 FSDP。该分片策略能够使单 GPU 的优化器内存开销随 $\text {GPU 个数}$ 线性降落，这就是为什么你会发现某个模型及其训练配置可能在大规模集群上能跑，但到小规模集群时就 OOM (Out Of Memory，内存耗尽) 了。下图来自于 ZeRO 论文，它形象地阐明了不同的 ZeRO 阶段及它们之间的区别 (留神 $P_{os}$、$P_{os+g}$ 和 $P_{os+g+p}$ 通常别离示意为 ZeRO-1、ZeRO-2、ZeRO-3。ZeRO-0 通常示意“禁用 ZeRO”):
上面，咱们总结一下 ZeRO 各阶段的内存开销公式 (假设咱们应用混合精度及 Adam 优化器):
- 对于 ZeRO-1，
$$
\begin {align*}\text {训练总内存} \approx \text {模型内存}+\frac {\text {优化器内存}}{\text {GPU 数}}+\text {激活内存}+\text {梯度内存}\end {align*}
$$
- 对于 ZeRO-2，
$$
\begin {align*}\text {训练总内存} \approx\text {模型内存}+\text {激活内存}+\frac {\text {优化器内存}+\text {梯度内存}}{\text {GPU 数}}\end {align*}
$$
- 对于 ZeRO-3，
$$
\begin {align*}\text {训练总内存} \approx \text {激活内存}+\frac {\text {模型内存}+\text {优化器内存}+\text {梯度内存}}{\text {GPU 数}} + \text {(ZeRO-3 实时参数量)}\end {align*}
$$
其中，在训练过程没有应用流水线并行或张量并行的条件下，$\text {GPU 数}$ 即为 $\text {DP 并行度}$。更多详细信息，请参阅 Sharded Optimizers + 3D Parallelism 一文。
请留神，ZeRO-3 引入了一组实时参数。这是因为 ZeRO-3 引入了一组配置项 (stage3\_max\_live\_parameters, stage3\_max\_reuse\_distance, stage3\_prefetch\_bucket\_size, stage3\_param\_persistence\_threshold) 来管制同一时刻 GPU 内存中能够放多少参数 (较大的值占内存更多但须要的通信更少)。这些参数会对总的 GPU 内存使用量产生重大影响。
请留神，ZeRO 还能够通过 ZeRO-R 在数据并行 rank 间划分激活，这样 $\text {激活内存}$ 还能够再除以张量并行度 $t$。更具体的信息，请参阅相干的 ZeRO 论文 及其 配置选项 (留神，在 GPT-NeoX 中，相应的配置标记为 partition_activations)。如果你正在训练一个大模型，激活放不下内存而成为一个瓶颈，你能够应用这个办法用通信换内存。把 ZeRO-R 与 ZeRO-1 联合应用时，内存耗费如下:
$$
\begin {align*}\text {训练总内存}\approx\text {模型内存}+\frac {\text {优化器内存}}{\text {GPU 数}}+\frac {\text {激活内存}}{\text {张量并行度}}+\text {梯度内存}\end {align*}
$$
3D 并行
LLM 次要有 3 种并行形式:
数据并行: 在多个模型正本间拆分数据
流水线或张量 / 模型并行: 在各 GPU 之间拆分模型参数，因而须要大量的通信开销。它们的内存开销大概是:
$$
\begin {align*}\text {并行后模型内存}\approx\frac {\text {模型内存}}{\text {流水线并行度}\times\text {张量并行度}}\end {align*}
$$
$$
\begin {align*}\text {并行后梯度内存}\approx\frac {\text {梯度内存}}{\text {流水线并行度}}\end {align*}
$$
请留神，这是个近似公式，因为 (1) 流水线并行对升高激活的内存需要没有帮忙、(2) 流水线并行要求所有 GPU 存储所有正在进行的 micro batch 的激活，这对大模型很重要、(3) GPU 须要长期存储并行计划所需的额定通信缓冲区。
分片优化器 + 3D 并行
当 ZeRO 与张量并行、流水线并行联合时，由此产生的 GPU 间的并行策略如下:
值得一提的是，数据并行度对于计算训练的全局 batch size 至关重要。数据并行度取决于你想在训练集群中放弃几份残缺模型正本:
$$
\begin {align*}\text {数据并行度 =}\frac {\text {GPU 数}}{\text {流水线并行度}\times\text {张量并行度}}\end {align*}
$$
尽管流水线并行和张量并行与所有 ZeRO 阶段都兼容 (例如，张量并行叠加上 ZeRO-3 后，咱们会首先对张量进行切片，而后在每个张量并行单元中利用 ZeRO-3)，但只有 ZeRO-1 与张量和 / 或流水线并行相结合时会成果才会好。这是因为梯度划分在不同并行策略间会有抵触 (如流水线并行和 ZeRO-2 都会对梯度进行划分)，这会显著减少通信开销。
把所有货色打包到一起，咱们能够失去一个典型的 3D 并行 + ZeRO-1 + 激活分区 计划:
$$
\begin {align*}\text {训练总内存} \approx\frac {\text {模型内存}}{\text {流水线并行度}\times\text {张量并行度}}+\frac {\text {优化器内存}}{\text {GPU 数}}+\frac {\text {激活内存}}{\text {张量并行度}}+\frac {\text {梯度内存}}{\text {流水线并行度}}\end {align*}
$$
总结
EleutherAI 的工程师常常应用上述估算办法来高效布局、调试分布式模型训练。咱们心愿廓清这些常常被忽视的但很有用的实现细节，如果你想与咱们探讨或认为咱们错过了一些好的办法，欢送你通过 mailto:contact@eleuther.ai 分割咱们！
请应用如下格局援用本文:
@misc {transformer-math-eleutherai,
title = {Transformer Math 101},
author = {Anthony, Quentin and Biderman, Stella and Schoelkopf, Hailey},
howpublished = \url {blog.eleuther.ai/},
year = {2023}
}
英文原文: https://blog.eleuther.ai/transformer-math/
原文作者: Quentin Anthony，Stella Biderman，Hailey Schoelkopf，作者均来自 EleutherAI
译者: Matrix Yao (姚伟峰)，英特尔深度学习工程师，工作方向为 transformer-family 模型在各模态数据上的利用及大规模模型的训练推理。
审校 / 排版: zhongdongy (阿东)