Transformer大模型在训练过程中所需的计算量

简介

许多关于Transformer语言模型的基本且重要的信息都可以用相当简单的方式计算出来。不幸的是，这些计算公式在NLP社区中并不广为人知。本文档的目的是收集这些公式，以及相关的知识，包括它们的来源和重要性。

**注意：**本文主要关注训练成本，而训练成本主要受VRAM的限制。有关推理成本的类似讨论，重点关注延迟，请查看Kipply撰写的这篇优秀的博客文章。

计算需求

计算训练Transformer模型成本的基本公式如下：

$C\approx \tau T=6PD$

其中：

• $C$ 是训练Transformer模型所需的总计算量，以浮点运算次数表示
• $C=C_{\text{forward}}+C_{\text{backward}}$
• $C_{\text{forward}}\approx 2PD$
• $C_{\text{backward}}\approx 4PD$
• $\tau$ 是硬件设置的总吞吐量（$\tau =(\text{GPU数量})\times (\text{实际FLOPs}/\text{GPU})$），以FLOPs为单位
• $T$ 是训练模型所花费的时间，以秒为单位
• $P$ 是Transformer模型中的参数数量
• $D$ 是数据集大小，以token为单位

这些公式在 OpenAI’s scaling laws paper and DeepMind’s scaling laws paper中提出并经过实验验证。有关更多信息，请参阅每篇论文。

值得一提的是，我们需要讨论一下 $C$ 的单位。$C$ 是总计算量的度量，但可以用许多单位来衡量，例如：

• FLOP-s，单位为 $[\frac{\text{浮点运算次数}}{\text{秒}}]\times [\text{秒}]$
• GPU-hours，单位为 $[\text{GPU数量}]\times [\text{小时}]$
• 缩放定律论文倾向于以PetaFLOP-days为单位报告值，即 $10^{15}\times 24\times 3600$ 次浮点运算

需要牢记的一个有用区别是 $\text{实际FLOPs}$ 的概念。虽然GPU加速器白皮书通常会宣传其理论FLOPs，但在实践中永远无法达到这些值（尤其是在分布式环境中！）。一些常见的分布式训练环境中报告的 $\text{实际FLOPs}$ 值在下面的计算成本部分报告。

请注意，参考了 这篇关于LLM训练成本的精彩博客文章 中使用的吞吐量-时间版本的成本公式。

参数与数据集的权衡

尽管严格来说，您可以根据需要训练任意数量token的Transformer，但训练的token数量会极大地影响计算成本和最终模型性能，因此在两者之间取得适当的平衡非常重要。

让我们先从最主要的问题开始：“计算最优”语言模型。通常被称为“Chinchilla缩放定律”，以论文中提出当前关于参数数量观点的模型系列命名，计算最优语言模型的参数数量和数据集大小满足近似值 $D=20P$。这在一个非常具体的意义上是最优的：在使用 1,000 个 GPU 1 小时和 1 个 GPU 1,000 小时的成本相同的资源机制下，如果您的目标是在最大化性能的同时最小化训练模型的 GPU 小时成本，则应使用上述公式。

**我们不建议训练token数量少于 200B 的 LLM。**尽管这对许多模型来说是“Chinchilla最优”的，但生成的模型通常质量很差。对于几乎所有应用，我们建议确定您的用例可接受的推理成本，并尽可能多地训练最大模型，以在该推理成本下尽可能多地训练token。

计算成本的工程意义

Transformer 的计算成本通常以 GPU-hours或者 FLOP-seconds为单位列出。

• GPT-NeoX 使用普通注意力机制在 A100 上实现了 150 TFLOP/s/A100 的吞吐量，使用 Flash 注意力机制则达到了 180 TFLOP/s/A100。这与其他高度优化的库在规模上是一致的，例如 Megatron-DS 报告在 A100 上的吞吐量在 137 到 163 TFLOP/s/A100 之间。
• 作为一般经验法则，您应该始终能够在 A100 上实现大约 120 TFLOP/s 的吞吐量。如果您看到低于 115 TFLOP/s/A100，则可能是您的模型或硬件配置有问题。
• 使用高质量的互连技术（如 InfiniBand），您可以在数据并行维度上实现线性或亚线性扩展（即，增加数据并行度应该会使整体吞吐量几乎线性增加）。下面是在橡树岭国家实验室的 Summit 超级计算机上测试 GPT-NeoX 库的图表。请注意，x 轴是 V100，而本文中的大多数数值示例都是针对 A100 的。

内存需求

Transformer 通常根据其参数大小来描述。但是，在确定哪些模型可以适合给定的计算资源集时，您需要知道模型将占用多少字节的空间。这可以告诉您，您的本地 GPU 可以容纳多大的模型进行推理，或者您可以在集群中使用一定数量的加速器内存训练多大的模型。

推理

模型权重

大多数 Transformer 都是以混合精度训练的，fp16 + fp32 或 bf16 + fp32。这减少了训练模型所需的内存量，也减少了运行推理所需的内存量。我们可以将语言模型从 fp32 转换为 fp16 甚至 int8，而不会造成明显的性能损失。这些数字指的是单个参数所需的位数(bits)。由于一个字节中有 8 位(bits)，我们将此数字除以 8，即可得出每个参数需要多少字节。

• 在 int8 中，${\text{memory}}_{\text{model}}=(1\text{ 字节}/\text{参数})\cdot (\text{参数数量})$
• 在 fp16 和 bf16 中， memory model = ( 2  字节 / 参数 ) ⋅ ( 参数数量 ) \text{memory}_{\text{model}}=(2 \text{ 字节} /\text{参数})\cdot ( \text{参数数量}) </