1
引言
Transformer 语言模型的许多根本重要信息能够通过简略计算得出。可怜的是，这些计算公式在自然语言解决（NLP）社区中并不广为人知。AI 非营利钻研组织 EleutherAI 收集整理这些公式，并介绍这些公式的起源和重要性。
注：本文次要关注显存（VRAM）主导的训练老本。无关推理老本和时延方面的相似探讨，请参见此前公布的《大型语言模型的推理演算》。
（本文经受权后由 OneFlow 编译公布，译文转载请分割 OneFlow 取得受权。原文：https://blog.eleuther.ai/transformer-math/）
2
计算需要
Transformer 模型的训练老本可通过以下根本公式计算得出：
这些公式由 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 是对总计算量的度量，能够用多种单位来示意。例如：
Actual FLOPs 是一个值得注意的有用概念。GPU 加速器白皮书通常会宣传它们的实践 FLOPs，但在理论状况下这些值往往无奈实现（特地是在分布式设置下）。上面的“计算成本”局部指出了在分布式训练设置中 Actual FLOPs 的常见报告值。
备注：咱们应用了吞吐量 - 工夫版本（throughput-time version）的老本公式，这个公式已在一篇对于 LLM 训练老本的文章中 (https://medium.com/@dzmitrybahdanau/the-flops-calculus-of-lan…) 应用。
参数与数据集衡量（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。
3
内存需要
通常，Transformer 模型的规模是以其参数数量为指标的。然而，在确定一个模型是否能够匹配特定计算资源时，咱们须要晓得模型将占用多少字节空间。这样能力弄清楚多大的模型适宜在本地 GPU 上进行推理，或者在肯定的总加速器内存下能够训练多大的模型。
推理
模型权重
大多数 Transformer 都是以混合精度训练的，不是 fp16 + fp32 就是 bf16 + fp32。这缩小了训练模型所需的内存量，以及运行推理所需的内存量。咱们能够将语言模型从 fp32 转换为 fp16 甚至 int8，而不会产生实质性的性能损失。这些数字指的是单个参数所需的位（bits）大小。因为一个字节（Byte）有 8 位，咱们把这个数除以 8，看看每个参数须要多少字节。
这里还有大量的额定开销，但通常与确定 GPU 所匹配的最大模型无关。依据教训来看，此额定开销≤20%。
总推理内存
除了存储模型权重所需的内存外，在理论前向传递过程中还会有大量额定开销。依据教训来看，此额定开销≤20%，并且通常与确定 GPU 所匹配的最大模型无关。
总的来说，可通过以下公式来“判断一个模型是否适宜进行推理”：
本文不会探讨这种额定开销的起源，这部分内容将会在其余文章中具体介绍。接下来将持续谈谈模型训练的内存。如果想理解更多对于推理所需的计算，请查看此前公布的《大型语言模型的推理演算》。
训练
除了模型参数之外，训练还须要在设施内存中存储优化器状态和梯度。这就是为什么当问及“我须要多少内存来拟合模型 X”时，会立刻失去“这取决于训练或推理”的答案。通常状况下，训练须要的内存比推理更多。
模型参数
首先，能够在纯 fp32 或 fp16 中训练模型：
‘’
除了推理中探讨的常见模型权重数据类型外，训练还引入了混合精度训练，例如 AMP。该技术旨在放弃收敛性的同时最大化 GPU 张量外围的吞吐量。古代深度学习训练畛域常常应用混合精度训练，因为：1) fp32 训练稳固，但内存额定开销高且不能利用 NVIDIA GPU 张量外围；2) fp16 训练稳固但难以收敛。留神：混合精度须要在内存中存储模型的 fp16/bf16 和 fp32 版本，因而须要：
再加上在优化器状态中计算的模型正本，大小为 (4 bytes/param) • (#params) 的正本。
优化器状态
Adam 很神奇，但它的内存效率非常低。除了要求领有模型参数和梯度参数的正本外，还须要额定保留三个梯度参数正本。所以，
梯度
梯度能够以 fp32 或 fp16 存储（留神，梯度数据类型通常与模型数据类型相匹配。因而，在应用 fp16 混合精度训练时，梯度通常以 fp16 存储），因而它们对内存额定开销的奉献为：
激活（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 模型激活所需内存的根本公式如下：
分布式训练
分片（Sharded）优化器
优化器的大量内存额定开销是 ZeRO 和 FSDP 等分片优化器的次要动机。这种分片策略将优化器额定开销升高了 No.GPUs 倍，这就是给定模型配置可能适宜大规模但 OOM 适宜小规模的起因。
在本博文中（假设为混合精度和 Adam 优化器）：
除非利用流水并行和 / 或张量并行，否则（DP Degree）就只是(No.GPUs)。相干详细信息，请参阅分片优化器 + 3D 并行局部
（https://eleutherai.notion.site/Transformers-Math-101-d2fcfc7a…）。
留神，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）上的激活进行分区。这也会使 在张量并行度 的下面。无关更多详细信息，请浏览相干 ZeRO 论文（https://arxiv.org/abs/1910.02054）和配置选项（https://www.deepspeed.ai/docs/config-json/#activation-checkpo…）（留神在 GPT-NeoX 中，这是 partition_activations 标识）。如果你正在训练巨型模型，想以一些内存额定开销换取额定的通信老本，那么激活会成为瓶颈问题。将 ZeRO- R 与 ZeRO- 1 一起应用的示例如下：
3D 并行
LLM 有 3 种次要并行形式：
数据并行：在模型正本（可能是模型并行的正本）之间划分（split）数据
流水或张量 / 模型并行：这些并行形式将模型的参数划分到 GPU 之间。此类计划须要大量通信的额定开销，但它们的内存的缩小大概为：
留神，此等式是近似值，起因如下：
（1）流水并行不会缩小激活的内存占用；
（2）流水并行须要所有 GPU 存储全副传输中的 micro-batches 激活，这对于大型模型十分重要；
（3）GPU 须要长期存储并行计划所需的额定通信缓冲区。
分片优化器 + 3D 并行
当 ZeRO 与张量和 / 或流水并行联合时，生成的并行策略会造成如下所示的网格：
4
总结
EleutherAI 的工程师们常常应用上述启发式办法布局高效的模型训练以及调试分布式运行。咱们心愿廓清这些常常被忽视的实现细节。
试用 OneFlow: github.com/Oneflow-Inc/oneflow/