Transformer 估算 101

本文主要介绍用于估算 transformer 类模型计算量需求和内存需求的相关数学方法。

引言

其实，很多有关 transformer 语言模型的一些基本且重要的信息都可以用很简单的方法估算出来。不幸的是，这些公式在 NLP 社区中鲜为人知。本文的目的是总结这些公式，阐明它们是如何推导出来的及其作用。

注意: 本文主要关注训练成本，该成本主要由 GPU 的 VRAM 主导。如果你想知道有关推理成本 (通常由推理延迟主导) 的信息，可以读读 Kipply 写的 这篇精彩博文。

算力要求

下式可用于估算训练一个 transformer 模型所需的算力成本:

这里:

• 是训练 transformer 模型所需的计算量，单位为总浮点运算数 (FLOP)

• 前向后向

• 前向

• 后向

• 是训练集群的实际总吞吐量: 数每的实际每秒浮点运算数实际，单位为 FLOPs

• 是训练模型所花费的时间，以秒为单位

• 是 transformer 模型的参数量

• 是数据集大小，表示为数据集的总词元数

该式由 OpenAI 的缩放定律论文 和 DeepMind 的缩放定律论文 提出并经其实验验证。想要获取更多信息，可参阅这两篇论文。

下面，我们讨论一下 的单位。 是总计算量的度量，我们可以用许多单位来度量它，例如:

• FLOP - 秒，单位为 每秒浮点运算数秒

• GPU - 时，单位为 数小时

• 缩放定律论文倾向于以 PetaFLOP - 天 为单位，即单位为

这里需要注意 实际 的概念。虽然 GPU 的规格书上通常只宣传其理论 FLOPs，但在实践中我们从未达到过这些理论值 (尤其在分布式训练时！)。我们将在计算成本这一小节列出分布式训练中常见的 实际 值。

请注意，上面的算力成本公式来自于 这篇关于 LLM 训练成本的精彩博文。

参数量与数据集的权衡

严格来讲，你可以随心所欲地使用任意数量的词元来训练 transformer 模型，但由于参与训练的词元数会极大地影响计算成本和最终模型的性能，因此需要小心权衡。

我们从最关键的部分开始谈起: “计算最优” 语言模型。“Chinchilla 缩放定律”，得名于提出 “计算最优” 语言模型论文中所训练的模型名，指出计算最优语言模型的 参数量 和 数据集大小 的近似关系满足: 。该关系成立基于一个前提条件: 使用 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</