作者 | Quentin Anthony、Stella Biderman、Hailey Schoelkopf

OneFlow编译

翻译 | 贾川、徐佳渝、杨婷

1
引言

Transformer语言模型的许多基本重要信息可以通过简单计算得出。不幸的是，这些计算公式在自然语言处理（NLP）社区中并不广为人知。AI非营利研究组织EleutherAI收集整理这些公式，并介绍这些公式的来源和重要性。

注：本文主要关注显存（VRAM）主导的训练成本。有关推理成本和时延方面的类似讨论，请参见此前发布的《大型语言模型的推理演算》。

（本文经授权后由OneFlow编译发布，译文转载请联系OneFlow获得授权。原文：https://blog.eleuther.ai/transformer-math/）

计算需求

Transformer模型的训练成本可通过以下基本公式计算得出：

C\approx\tau T = 6PD

其中：

C表示训练Transformer模型所需的计算量，以总浮点运算数为单位。
$C=C_{\text{forward}}+C_{\text{backward}}$
$C_{\text{forward}}\approx2PD$
$C_{\text{backward}}\approx4PD$
$\tau$ 表示硬件设置的总吞吐量（ $\tau$ =(No.GPUs)x（Actual FLOPs/GPU)），以FLOPs为单位。
T表示训练模型所需的时间，以秒为单位。
P表示Transformer模型中的参数数量。
D表示数据集大小，以token为单位。

这些公式由OpenAI的论文《Scaling Laws for Neural Language Models》(https://arxiv.org/abs/2001.08361)和DeepMind的论文《Training Compute-Optimal Large Language Models》(https://arxiv.org/abs/2203.15556)提出并经实验证实。

值得一提的是C的单位问题。C是对总计算量的度量，可以用多种单位来表示。例如：

FLOP-seconds，单位为 $[\frac{\text{Floating Point Operations}}{\text{Second}}] \times [\text{Seconds}]$
GPU-hours，单位为[No.GPUs] x [Hours]
扩展定律（Scaling laws）论文通常使用PetaFLOP-days，或者 $10^{15}\times24\times3600$ 总浮点运算数来报告值（values）。

Actual FLOPs是一个值得注意的有用概念。GPU加速器白皮书通常会宣传它们的理论FLOPs，但在实际情况下这些值往往无法实现（特别是在分布式设置下）。下面的“计算成本”部分指出了在分布式训练设置中Actual FLOPs的常见报告值。

备注：我们使用了吞吐量-时间版本（throughput-time version）的成本公式，这个公式已在一篇关于LLM训练成本的文章中(https://medium.com/@dzmitrybahdanau/the-flops-calculus-of-language-model-training-3b19c1f025e4)使用。

参数与数据集权衡（Tradeoffs）

严格来说，你可以使用任意数量的token训练Transformer模型。然而，训练的token数量会对计算成本和最终模型的性能产生很大影响，因此找到一个合适的平衡点非常重要。

让我们从关键问题“计算最优（compute optimal）”语言模型入手。 关于参数数量的理论“Chinchilla scaling laws”指出，一个计算最优的语言模型拥有的参数数量和数据集大小需满足一个近似等式：D=20P。

这在特定意义上是最优的：假设使用1000个GPU运行1小时的成本与使用1个GPU运行1000小时的成本相同，如果你的目标是在最小化GPU训练成本的同时最大化模型性能，你就应该使用上面提到的公式。

我们不建议训练小于200B个token的大型语言模型。 虽然对于许多模型来说这样才能达到“Chinchilla最优”，但最终的模型效果通常非常差。针对所有应用程序，我们建议要确定用例的可接受推理成本，然后尽可能地训练最大的模型，以使得尽可能多的token的训练保持在这一推理成本之下。

计算成本的工程要点

Transformer模型通常用GPU-hours或FLOP-seconds来表示计算成本。

GPT-NeoX在Normal Attention下实现了150 TFLOP/s/A100，而在Flash Attention下实现了180 FLOP/s/A100。这与其他高度优化的库在规模上是相符的，例如Megatron-DS的报告是在137至163 TFLOP/s/A100之间。
就经验来看，我们应该始终能够达到120 FLOP/s/A100的计算性能。如果计算性能低于115 FLOP/s/A100，那么很可能是模型或硬件配置出了问题。
使用高质量的互连技术（例如 InfiniBand），我们可以在数据并行维度上实现线性或亚线性扩展（sublinear scaling，即增加数据并行度应该近乎线性地增加总吞吐量）。下图展示了在Oak Ridge National Lab的Summit超级计算机上测试GPT-NeoX库的结果。注意：x 轴上是 V100，但文章中大部分数值示例针对的都是A100。

内存需求

通常，Transformer模型的规模是以其参数数量为指标的。然而，在确定一个模型是否可以匹配特定计算资源时，我们需要知道模型将占用多少字节空间。这样才能弄清楚多大的模型适合在本地GPU上进行推理，或者在一定的总加速器内存下可以训练多大的模型。

推理

模型权重

大多数Transformer都是以混合精度训练的，不是fp16 + fp32就是 bf16 + fp32。这减少了训练模型所需的内存量，以及运行推理所需的内存量。我们可以将语言模型从fp32转换为fp16甚至int8，而不会产生实质性的性能损失。这些数字指的是单个参数所需的位（bits）大小。由于一个字节（Byte）有8位，我们把这个数除以8，看看每个参数需要多少字节

int8， $\text{memory}_{\text{model}}=(1 \text{ byte} /\text{param})\cdot ( \text{No. params})$
fp16 and bf16， $\text{memory}_{\text{model}}=(2 \text{ bytes} /\text{param})\cdot ( \text{No. params})$
fp32， $\text{memory}_{\text{model}}=(4 \text{ bytes} /\text{param})\cdot (\text{No. params})$

这里还有少量的额外开销，但通常与确定GPU所匹配的最大模型无关。根据经验来看，此额外开销≤20%。

总推理内存

除了存储模型权重所需的内存外，在实际前向传递过程中还会有少量额外开销。根据经验来看，此额外开销≤20%，并且通常与确定GPU所匹配的最大模型无关。

总的来说，可通过以下公式来“判断一个模型是否适合进行推理”：

\text{Total Memory}_{\text{Inference}}\approx(1.2) \times \text{Model Memory}

本文不会探讨这种额外开销的来源，这部分内容将会在其他文章中详细介绍。接下来将继续谈谈模型训练的内存。如果想了解更多关于推理所需的计算，请查看此前发布的《大型语言模型的推理演算》。

训练

除了模型参数之外，训练还需要在设备内存中存储优化器状态和梯度。这就是为什么当问及“我需要多少内存来拟合模型X”时，会立即得到“这取决于训练或推理”的答案。通常情况下，训练需要的内存比推理更多。

模型参数

首先，可以在纯fp32或fp16中训练模型：

纯fp32， $\text{memory}_{\text{model}}=(4 \text{ bytes} /\text{param})\cdot (\text{No. params})$
纯fp16， $\text{memory}_{\text{model}}=(2 \text{ bytes} /\text{param})\cdot (\text{No. params})$

除了推理中讨论的常见模型权重数据类型外，训练还引入了混合精度训练，例如AMP。该技术旨在保持收敛性的同时最大化GPU张量核心的吞吐量。现代深度学习训练领域经常使用混合精度训练，因为：1) fp32训练稳定，但内存额外开销高且不能利用NVIDIA GPU张量核心；2) fp16训练稳定但难以收敛。注意：混合精度需要在内存中存储模型的fp16/bf16和fp32版本，因此需要：

混合精度（fp16/bf16 和 fp32）， $\text{memory}_{\text{model}}=(2 \text{ bytes} /\text{param})\cdot (\text{No. params})$

再加上在优化器状态中计算的模型副本，大小为(4 bytes/param) • (#params)的副本。

优化器状态

Adam很神奇，但它的内存效率非常低。除了要求拥有模型参数和梯度参数的副本外，还需要额外保留三个梯度参数副本。所以，

对于vanilla AdamW， $\text{memory}_{\text{optimizer}}=(12 \text{ bytes}/\text{param})\cdot (\text{No. params})$

fp32参数副本：4字节/参数
动量：4字节/参数
方差：4字节/参数

对于像bitsandbytes这样的8位优化器，

$\text{memory}_{\text{optimizer}}=(6 \text{ bytes} /\text{param})\cdot (\text{No. params})$

fp32参数副本：4字节/参数
动量：1字节/参数
方差：1字节/参数

对于具有动量的类SGD优化器， $\text{memory}_{\text{optimizer}}=(8 \text{ bytes} /\text{param})\cdot (\text{No. params})$

fp32参数副本：4字节/参数
动量：4字节/参数

梯度

梯度可以以fp32或fp16存储（注意，梯度数据类型通常与模型数据类型相匹配。因此，在使用fp16混合精度训练时，梯度通常以fp16存储），因此它们对内存额外开销的贡献为：

fp32， $\text{memory}_{\text{gradients}}=(4 \text{ bytes} /\text{param})\cdot (\text{No. params})$
fp16， $\text{memory}_{\text{gradients}}=(2 \text{ bytes} /\text{param})\cdot (\text{No. params})$

激活（Activations）和批量大小

对于大型语言模型的训练来说，现代GPU通常受制于内存而非浮点运算。因此，激活重计算（activation recomputation又称checkpointing）是一种非常流行的方法，可以用减少的内存成本换取额外的计算成本。

激活重计算的工作原理是重新计算某些层的激活，而不是将它们存储在GPU内存中。内存的减少取决于我们在决定清除哪些层时所作的选择，但Megatron的选择性重计算方案如下图所示：

其中，红色虚线表示 A100-80GB GPU 的内存容量，“present work”表示应用选择性激活重计算后的内存需求。更多有关信息和下列公式的推导，请参阅Reducing Activation Recomputation in Large Transformer Models

（https://arxiv.org/abs/2205.05198）。

存储Transformer模型激活所需内存的基本公式如下：

\begin{align*}\text{memory}^{\text{No Recomputation}}_{\text{activations}}=sbhL(10+\frac{24}{t}+5\frac{a \cdot s}{h\cdot t}) \text{ bytes}\end{align*}

\begin{align*}\text{memory}^{\text{Selective Recomputation}}_{\text{activations}}=sbhL(10+\frac{24}{t}) \text{ bytes}\end{align*}

\begin{align*}\text{memory}^{\text{Full Recomputation}}_{\text{activations}}=2 \cdot sbhL \text{ bytes}\end{align*}

其中：

$s$ 是序列长度，以token为单位
$b$ 是每个GPU的批量大小
$h$ 是每个Transformer层中hidden size的维度
$L$ 是Transformer模型中的层数
$a$ 是Transformer模型中attention heads的数量
$t$ 是使用的张量并行度（如果不是则为 1）
假设没有使用序列并行
假设激活存储在fp16中

此外，所需的额外重计算也取决于方法的选择，但它受限于完整的额外前向传递。因此，前向传递的更新成本由以下公式给出：

2PD\leq C_{\text{forward}}\leq4PD

总训练内存

因此，有关“模型是否适合训练”的较好的启发式答案如下：

\begin{align*}\text{Total Memory}_{\text{Training}} = \text{Model Memory}+\text{Optimiser Memory}+\text{Activation Memory}+\text{Gradient Memory}\end{align*}

分布式训练

分片（Sharded）优化器

优化器的大量内存额外开销是ZeRO和FSDP等分片优化器的主要动机。这种分片策略将优化器额外开销降低了No.GPUs 倍，这就是给定模型配置可能适合大规模但OOM适合小规模的原因。

如果要计算使用分片优化器进行训练所需的内存额外开销，则需要包含下图中的公式。有关分片优化的示例计算，请参见ZeRO论文中的下图（注意，

P_{os}P_{os+g}和 P_{os+g+p}

通常分别表示为 ZeRO-1、ZeRO-2、ZeRO-3。ZeRO-0 通常意为“ZeRO 禁用”）：

在本博文中（假定为混合精度和Adam优化器）：

ZeRO-1，

\begin{align*}\text{Total Memory}_{\text{Training}}\approx\text{Model Memory}+\frac{\text{Optimizer memory}}{(\text{No. GPUs})}+\text{Activation Memory}+\text{Gradient Memory}\end{align*}

ZeRO-2，

\begin{align*}\text{Total Memory}_{\text{Training}}\approx\text{Model Memory}+\text{Activation Memory}+\frac{\text{Optimizer Memory}+\text{Gradient Memory}}{(\text{No. GPUs})}\end{align*}

ZeRO-3，

\begin{align*}\text{Total Memory}_{\text{Training}}\approx \text{Activation Memory}+\frac{\text{Model Memory}+\text{Optimizer Memory}+\text{Gradient Memory}}{(\text{No. GPUs})} + \text{(ZeRO-3 Live Params)}\end{align*}

除非应用流水并行和/或张量并行，否则（DP Degree）就只是(No.GPUs)。相关详细信息，请参阅分片优化器 + 3D并行部分

（ https://eleutherai.notion.site/Transformers-Math-101-d2fcfc7a25d446388fde97821ad2412a#9c476d020d7641a299fb6be6ae82e9f8）。

注意，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对数据并行列（ranks）上的激活进行分区。这也会使

\text{memory}_\text{activations}

在张量并行度

t

的上面。有关更多详细信息，请阅读相关ZeRO论文（https://arxiv.org/abs/1910.02054）和配置选项（https://www.deepspeed.ai/docs/config-json/#activation-checkpointing）（注意在GPT-NeoX中，这是partition_activations标识）。如果你正在训练巨型模型，想以一些内存额外开销换取额外的通信成本，那么激活会成为瓶颈问题。将ZeRO-R与ZeRO-1一起使用的示例如下：

\begin{align*}\text{Total Memory}_{\text{Training}}\approx\text{Model Memory}+\frac{\text{Optimizer Memory}}{(\text{No. GPUs})}+\frac{\text{Activation Memory}}{\text{(Tensor-Parallel-Size)}}+\text{Gradient Memory}\end{align*}

3D并行

LLM有3种主要并行方式：

数据并行： 在模型副本（可能是模型并行的副本）之间划分（split）数据

流水或张量/模型并行： 这些并行方式将模型的参数划分到GPU之间。此类方案需要大量通信的额外开销，但它们的内存的减少大约为：

\begin{align*}\text{memory}^{\text{w/ parallelism}}_{\text{model}}\approx\frac{\text{Model Memory}}{\text{(Pipe-Parallel-Size})\times\text{(Tensor-Parallel-Size)}}\end{align*}

\begin{align*}\text{memory}^{\text{w/ parallelism}}_{\text{gradients}}\approx\frac{\text{Gradient Memory}}{\text{(Pipe-Parallel-Size})}\end{align*}

注意，此等式是近似值，原因如下：

（1）流水并行不会减少激活的内存占用；

（2）流水并行需要所有GPU存储全部传输中的micro-batches激活，这对于大型模型非常重要；

（3）GPU需要临时存储并行方案所需的额外通信缓冲区。

分片优化器 + 3D 并行

当ZeRO与张量和/或流水并行结合时，生成的并行策略会形成如下所示的网格：

一个重要的题外话：数据并行度（DP degree)对于计算训练的全局批量大小至关重要。数据并行度取决于完整模型副本数量：

\begin{align*}\text{DP Degree = }\frac{\text{No. GPUs}}{\text{(Pipe-Parallel-Size})\times\text{(Tensor-Parallel-Size)}}\end{align*}

虽然流水并行和张量并行与ZeRO的所有阶段兼容（例如，张量并行的ZeRO-3会先对张量切片，然后在每个张量并行单元中应用ZeRO-3），但只有ZeRO-1倾向于与张量和/或流水并行结合使用时表现良好，这是因为梯度并行策略相互冲突（流水并行和ZeRO-2是划分梯度）（张量并行和ZeRO-3是划分模型参数），从而造成大量通信的额外开销。

将所有内容放在一起，进行典型的带有激活分区的3D并行ZeRO-1运行：

\begin{align*}\text{Total Memory}_{\text{Training}}\approx\frac{\text{Model Memory}}{\text{(Pipe-Parallel-Size})\times\text{(Tensor-Parallel-Size)}}+\frac{\text{Optimizer Memory}}{(\text{No. GPUs})}+\frac{\text{Activation Memory}}{\text{(Tensor-Parallel-Size)}}+\frac{\text{Gradient Memory}}{\text{(Pipe-Parallel-Size})}\end{align*}

总结

EleutherAI的工程师们经常使用上述启发式方法规划高效的模型训练以及调试分布式运行。我们希望澄清这些经常被忽视的实现细节。

其他人都在看

大型语言模型的推理演算

向量嵌入：AutoGPT的幻觉解法

John Schulman：通往TruthGPT之路

复杂推理：大型语言模型的"北极星"能力

为什么ChatGPT用强化学习而非监督学习

OneEmbedding:单卡 训练TB级推荐模型不是梦

GLM训练加速：性能最高提升3倍，显存节省1/3

试用OneFlow: github.com/Oneflow-Inc/oneflow/

本文分享自微信公众号 - OneFlow（OneFlowTechnology）。
如有侵权，请联系 support@oschina.cn 删除。
本文参与“OSC源创计划”，欢迎正在阅读的你也加入，一起分享。

魔力猫 2023-11-30 10:23

如果老版本照旧，新版本文档封闭，这没太大问题。遗留系统，各种烂事多了去了。升级框架要加钱不算啥。可如果把老文档也封了，就等同绑票！而且作者也别吹已经有多少肉票交了赎金。交钱的有多少心里咒你们没点数？如果不是沉默成本太大，谁愿意给这笔钱？！考虑过这个框架但没用的，那是真庆幸自己没被杀猪。以后不是真正大厂支持的项目，敢随便放到公司里面了。像这次的事情，风险没边了！今天封框架文档，明天就可以封部分API，不给钱不能用。MIT协议在眼里不如一张厕纸。如果程序员把框架推荐引入公司项目，结果出了这么一出，尴尬不？接下来怎么办？跟公司说现在要交赎金，遇到苛刻的老板，为这499就让你拍屁股走人，还不给补偿，罪名都是现成的，给公司造成重大损失，499不是钱的事情，而是项目风险不可控。

单一结构 2023-12-02 16:18

注意“允许任何人”。而他是“任意任何给与费用的人”，存在本质区别。别洗了

百小僧 2023-11-30 20:30

👉👉👉 说句实话，有时候我还是想不通你们这些人的。我们认识吗？我们有那么大的仇恨吗？我是伤害了你的利益吗？你了解的事情经过吗？你听到的看到的都是真的吗？你怀疑过这一切吗？你有限的知识够评价别人无限的人生吗？互联网那些被网暴的人真的就是道德败坏的人吗？你看懂社会运行规则了吗？人性的弱点你读懂了吗？你躬身入局了吗？自然辩证法学过吗？王阳明的心学心外无物课外阅读过吗？捕风捉影、愤然怒喷，把自己放在上帝的视角藐视着这一切，彷佛是一切的主宰，道德的化身，法律的审判者，人间的正义使者。

iro 2023-12-02 15:47

挂得好

技术粉 2023-12-02 19:42

不要以结果推条件，这是不对的

Transformer模型的基础演算

1
引言