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/