PaLM：使用 Pathways 进行语言建模的扩展

Aakanksha Chowdhery*, Sharan Narang*, Jacob Devlin* Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh, Kensen Shi Sasha Tsvyashchenko, Joshua Maynez, Abhishek Rao†, Parker Barnes, Yi Tay Noam Shazeer‡, Vinodkumar Prabhakaran, Emily Reif, Nan Du, Ben Hutchinson Reiner Pope, James Bradbury, Jacob Austin, Michael Isard, Guy Gur-Ari Pengcheng Yin, Toju Duke, Anselm Levskaya, Sanjay Ghemawat, Sunipa Dev Henryk Michalewski, Xavier Garcia, Vedant Misra, Kevin Robinson, Liam Fedus Denny Zhou, Daphne Ippolito, David Luan‡, Hyeontaek Lim, Barret Zoph Alexander Spiridonov, Ryan Sepassi, David Dohan, Shivani Agrawal, Mark Omernick Andrew M. Dai, Thanumalayan Sankaranarayana Pillai, Marie Pellat, Aitor Lewkowycz Erica Moreira, Rewon Child, Oleksandr Polozov†, Katherine Lee, Zongwei Zhou Xuezhi Wang, Brennan Saeta, Mark Diaz, Orhan Firat, Michele Catasta†, Jason Wei Kathy Meier-Hellstern, Douglas Eck, Jeff Dean, Slav Petrov, Noah Fiedel

Google Research

---

摘要

大型语言模型已被证明可以通过少样本学习（few-shot learning）在各种自然语言任务中取得显著的性能，这极大地减少了将模型适配到特定应用所需的任务特定训练示例数量。为了进一步加深我们对规模对少样本学习影响的理解，我们训练了一个 5400 亿参数、密集激活的 Transformer 语言模型，我们称之为 Pathways 语言模型（PaLM）。

我们使用 Pathways 在 6144 个 TPU v4 芯片上训练了 PaLM，Pathways 是一种新的机器学习系统，能够跨多个 TPU Pod 进行高效训练。我们通过在数百个语言理解和生成基准测试中取得最先进的少样本学习结果，证明了扩展带来的持续收益。在其中许多任务上，PaLM 540B 实现了突破性的性能，在一系列多步推理任务中超越了微调后的最先进水平，并在最近发布的 BIG-bench 基准测试中超越了人类的平均表现。相当多的 BIG-bench 任务显示出模型规模带来的不连续改进，这意味着随着我们扩展到最大的模型，性能得到了大幅提升。PaLM 在多语言任务和源代码生成方面也具有强大的能力，我们在广泛的基准测试中证明了这一点。此外，我们还对偏见和毒性进行了全面分析，并研究了训练数据记忆与模型规模的关系。最后，我们讨论了与大型语言模型相关的伦理考量，并探讨了潜在的缓解策略。

---

1. 引言

在过去的几年里，为语言理解和生成而训练的超大规模神经网络在各种任务中取得了惊人的成果。许多此类模型，如 BERT (Devlin et al., 2019) 和 T5 (Raffel et al., 2020)，都是使用仅编码器（encoder-only）或编码器-解码器（encoder-decoder）架构，通过在大型文本语料库上使用填充（“掩码语言模型”或“跨度损坏”）预训练目标进行训练，通常随后通过微调适配到一个或多个特定任务。尽管这些模型在数千个自然语言任务中几乎达到了普遍的最先进水平，但其缺点是需要大量的任务特定训练示例来微调模型。此外，至少部分模型参数必须更新以适应任务，这增加了模型微调和部署的复杂性。

GPT-3 (Brown et al., 2020) 证明了超大规模自回归语言模型（LMs）可用于少样本预测，其中模型仅被赋予自然语言任务描述和（可选的）少量演示任务应如何完成的示例。这类模型使用仅解码器（decoder-only）架构和标准的从左到右语言建模目标在大型文本语料库上进行训练，其目标是根据示例中的前序标记预测下一个标记。少样本评估已被证明可以取得非常强大的结果，而无需大规模的任务特定数据收集或模型参数更新。

自 GPT-3 以来，已经开发了许多其他大型自回归语言模型，它们继续推动着最先进技术的发展。这些后 GPT-3 模型中最强大的是 GLaM (Du et al., 2021)、Gopher (Rae et al., 2021)、Chinchilla (Hoffmann et al., 2022)、Megatron–Turing NLG (Smith et al., 2022) 和 LaMDA (Thoppilan et al., 2022)，它们在发布时都在大量任务上取得了少样本最先进的结果。与 GPT-3 一样，这些模型都是 Transformer 架构 (Vaswani et al., 2017) 的变体。这些模型的改进主要来自以下一种或多种方法：（1）在深度和宽度上扩展模型规模；（2）增加模型训练所用的标记（token）数量；（3）在来自更多样化来源的更干净数据集上进行训练；以及（4）通过稀疏激活模块在不增加计算成本的情况下提高模型容量。

在这项工作中，我们延续了语言建模改进的扩展路线，并在 7800 亿个高质量文本标记上训练了一个 5400 亿参数、密集激活的自回归 Transformer。这是通过使用 Pathways (Barham et al., 2022) 实现的，这是一种新的机器学习系统，能够跨数千个加速器芯片（包括跨越多个张量处理单元 (TPU) v4 Pod 的芯片）高效训练超大规模神经网络。这个名为 Pathways 语言模型（PaLM）的新模型在数百个自然语言、代码和数学推理任务中实现了最先进的少样本结果。我们在其中许多任务上实现了突破性的性能。

本工作的主要结论如下：

• 高效扩展 – 我们展示了 Pathways (Barham et al., 2022) 的首次大规模使用——这是一种新的机器学习系统，能够以极高效率跨数千或数万个加速器芯片训练单个模型。借助 Pathways，我们在 6144 个 TPU v4 芯片上训练了一个 540B 参数的语言模型，其效率水平是此类规模模型此前无法达到的。大多数以前的大型语言模型要么是在单个 TPU 系统上训练的 (Du et al., 2021; Thoppilan et al., 2022)，要么是使用流水线并行 (Huang et al., 2019) 跨 GPU 集群 (Smith et al., 2022) 或多个 TPU v3 Pod (Rae et al., 2021) 进行扩展，最大规模为 4096 个 TPU v3 芯片。在第 4 节中，我们描述了如何将 PaLM 540B 的无流水线训练扩展到跨两个 TPU v4 Pod 的 6144 个芯片，同时实现了 46.2% 的模型 FLOPs 利用率（观察到的吞吐量相对于理论最大吞吐量）和 57.8% 的硬件 FLOPs 利用率。
• 扩展带来的持续改进 – 在第 6 节中，我们在数百个自然语言、代码和数学推理任务上评估了 PaLM，并在绝大多数基准测试中取得了最先进的结果，通常优势显著。这关键地证明了大型语言模型的扩展改进既没有达到平台期，也没有达到饱和点。例如，在表 4 中，我们展示了在 29 个最广泛评估的英语语言理解基准测试中的 28 个上，新的少样本最先进水平，与 GLaM (Du et al., 2021)、GPT-3 (Brown et al., 2020)、Megatron–Turing NLG (Smith et al., 2022)、Gopher (Rae et al., 2021)、Chinchilla (Hoffmann et al., 2022) 和 LaMDA (Thoppilan et al., 2022) 的最佳单任务结果相比。

• 突破性能力 – 我们展示了在许多困难任务中语言理解和生成的突破性能力。具体而言，第 6.3 节介绍了对一系列推理任务的评估，这些任务需要多步数学或常识推理才能得出正确答案。先前的最先进结果使用任务特定微调、领域特定架构和任务特定验证器的组合来取得强大结果。在这项工作中，我们证明了当模型扩展与思维链提示 (Wei et al., 2022b) 相结合时，简单的少样本评估可以在广泛的推理任务上超越或匹配微调后的最先进水平。在第 6.2 节中，我们还重点介绍了在 BIG-bench (BIG-bench collaboration, 2021) 上的突破性性能，这是一套最近发布的 150 多个新语言理解和生成任务，其中许多任务即使对人类来说也非常难以正确回答。在图 1 和第 9 节中，我们展示了 PaLM 在显式解释和说明复杂推理链方面的探索性能力。
• 不连续的改进 – 为了更好地理解扩展行为，我们展示了三种参数规模的结果：8B、62B 和 540B。通常，从 62B 扩展到 540B 的性能提升与从 8B 扩展到 62B 的性能提升相似，这与神经网络扩展中经常观察到的“幂律”经验法则一致 (Kaplan et al., 2020)。然而，对于某些任务，我们观察到不连续的改进，即从 62B 扩展到 540B 导致的准确率跳跃比从 8B 扩展到 62B 要剧烈得多。这种行为在第 6.2 节中约 25% 的 BIG-bench 任务中观察到。这表明当模型达到足够的规模时，大型语言模型的新能力可能会涌现，并且这些能力在超过先前研究的规模后会继续涌现。
• 多语言理解 – 以前关于大型语言模型的工作在多语言领域进行了有限的评估。在这项工作中，我们对多语言基准测试进行了更彻底的评估，包括机器翻译（第 6.5 节）、摘要（第 6.6 节）和问答（第 6.7 节），涵盖了多种语言。即使训练语料库中非英语数据的比例相对较小（≈ 22%），540B 模型的少样本评估结果也能够弥合与非英语摘要任务中先前微调的最先进水平之间的差距，并在翻译任务中超越先前最先进水平。需要进一步的工作来了解增加多语言数据比例对英语和多语言任务的影响。
• 偏见和毒性 – 我们还评估了模型在分布偏见和毒性方面的表现，这产生了一些见解（第 10 节）。首先，对于性别和职业偏见，我们发现 Winogender 指代消解任务的准确率随着模型规模的增加而提高，PaLM 540B 在 1-shot 和少样本设置中设定了新的最先进结果。其次，对种族/宗教/性别提示续写进行的共现分析表明，该模型有可能错误地确认刻板印象，例如将穆斯林与恐怖主义、极端主义和暴力联系起来。这种行为在不同模型规模中是一致的。最后，对提示续写任务的毒性分析表明，与 8B 模型相比，62B 和 540B 模型的总体毒性水平略高。然而，模型生成的续写的毒性与提示文本的毒性高度相关，而人类生成的续写则没有很强的毒性相关性。这表明模型受提示风格的影响比人类生成的文本更大。在未来的工作中，我们计划将这些基准测试扩展到非英语语言，并更详尽地考虑潜在风险。

---

2. 模型架构

PaLM 使用标准的 Transformer 模型架构 (Vaswani et al., 2017)，采用仅解码器设置（即每个时间步只能关注自身和过去的时间步），并进行了以下修改：

• SwiGLU 激活 – 我们使用 SwiGLU 激活（$Swish(xW) \cdot xV$）作为 MLP 中间激活，因为已证明与标准的 ReLU、GeLU 或 Swish 激活相比，它们能显著提高质量 (Shazeer, 2020)。请注意，这确实需要在 MLP 中进行三次矩阵乘法而不是两次，但 Shazeer (2020) 在计算等效实验中证明了质量的提升（即标准 ReLU 变体具有成比例更大的维度）。
• 并行层 – 我们在每个 Transformer 块中使用“并行”公式 (Wang & Komatsuzaki, 2021)，而不是标准的“串行”公式。具体而言，标准公式可以写为： $y = x + \text{MLP}(\text{LayerNorm}(x + \text{Attention}(\text{LayerNorm}(x))))$ 而并行公式可以写为： $y = x + \text{MLP}(\text{LayerNorm}(x)) + \text{Attention}(\text{LayerNorm}(x))$ 并行公式在大规模下训练速度快约 15%，因为 MLP 和 Attention 输入矩阵乘法可以融合。消融实验显示 8B 规模下质量有轻微下降，但在 62B 规模下没有质量下降，因此我们推断并行层在 540B 规模下应该是质量中性的。
• 多查询注意力（Multi-Query Attention） – 标准 Transformer 公式使用 $k$ 个注意力头，其中每个时间步的输入向量被线性投影为形状为 $[k, h]$ 的“查询”、“键”和“值”张量，其中 $h$ 是注意力头大小。在这里，键/值投影在每个头之间共享，即“键”和“值”被投影到 $[1, h]$，但“查询”仍然被投影到形状 $[k, h]$。我们发现这对模型质量和训练速度有中性影响 (Shazeer, 2019)，但在自回归解码时显著节省了成本。这是因为标准的多头注意力在自回归解码期间在加速器硬件上的效率较低，因为键/值张量不在示例之间共享，并且一次只解码一个标记。
• RoPE 嵌入 – 我们使用 RoPE 嵌入 (Su et al., 2021) 而不是绝对或相对位置嵌入，因为已证明 RoPE 嵌入在长序列长度上具有更好的性能。
• 共享输入-输出嵌入 – 我们共享输入和输出嵌入矩阵，这在过去的工作中经常（但不普遍）这样做。
• 无偏置 – 在任何密集核或层归一化中都没有使用偏置。我们发现这可以提高大型模型的训练稳定性。
• 词汇表 – 我们使用具有 256k 个标记的 SentencePiece (Kudo & Richardson, 2018a) 词汇表，选择它是为了在不进行过度分词的情况下支持训练语料库中的大量语言。词汇表是从训练数据中生成的，我们发现这提高了训练效率。词汇表是完全无损且可逆的，这意味着空格在词汇表中被完全保留（这对代码特别重要），并且词汇表外的 Unicode 字符被拆分为 UTF-8 字节，每个字节都有一个词汇表标记。数字总是被拆分为单独的数字标记（例如，“123.5 → 1 2 3 . 5”）。

2.1 模型规模超参数

在这项工作中，我们比较了三种不同的模型规模：540B 参数、62B 参数和 8B 参数。每个标记的 FLOPs 数量大约等于参数数量，因为这些模型是标准的密集 Transformer。这些模型是使用表 1 中的超参数构建的。这三个模型使用相同的数据和词汇表进行了相同的训练（批量大小除外）。训练在第 3 节和第 5 节中有更详细的描述。

模型	层数	头数	$d_{model}$	参数数量（十亿）	批量大小
PaLM 8B	32	16	4096	8.63	256 → 512
PaLM 62B	64	32	8192	62.50	512 → 1024
PaLM 540B	118	48	18432	540.35	512 → 1024 → 2048

表 1：模型架构详情。我们列出了层数、$d_{model}$、注意力头数和注意力头大小。前馈大小 $d_{ff}$ 始终为 $4 \times d_{model}$，注意力头大小始终为 256。

2.2 模型卡

PaLM 的模型卡 (Mitchell et al., 2019) 在附录 E 中提供。这提供了模型架构、训练设置、训练数据和预期用途的高级摘要。

---

3. 训练数据集

PaLM 预训练数据集由 7800 亿个高质量标记组成，代表了广泛的自然语言用例。该数据集是过滤后的网页、书籍、Wikipedia、新闻文章、源代码和社交媒体对话的混合体。该数据集基于用于训练 LaMDA (Thoppilan et al., 2022) 和 GLaM (Du et al., 2021) 的数据集。我们对所有三个模型进行完全一轮（epoch）的数据训练（所有模型的数据洗牌方式相同），并选择混合比例以避免在任何子组件中重复数据。

除了自然语言数据外，预训练数据集还包含代码。预训练数据集中的源代码是从 GitHub 上的开源存储库中获取的。我们根据存储库中包含的许可证过滤了文件；排除了 copyleft 许可证。我们通过文件名扩展名过滤文件，以限制为 24 种常见编程语言之一，包括 Java、HTML、Javascript、Python、PHP、C#、XML、C++ 和 C，这产生了 196GB 的源代码。此外，我们根据文件之间的 Levenshtein 距离删除了重复项，因为已知重复文件在源代码存储库中很常见 (Lopes et al., 2017; Allamanis, 2019)。

表 2 列出了用于创建最终 PaLM 数据集混合的各种数据源的比例。我们在第 8 节中检查了数据污染，并展示了我们的训练数据集与评估数据之间的重叠分析。附录 D 包含一份数据表 (Gebru et al., 2021)，其中包含更多信息，包括语言比例的细分。

数据源	数据比例
社交媒体对话（多语言）	50%
过滤后的网页（多语言）	27%
书籍（英语）	13%
GitHub（代码）	5%
Wikipedia（多语言）	4%
新闻（英语）	1%

表 2：训练数据集中每个来源的数据比例。多语言语料库包含来自 100 多种语言的文本，分布情况见附录表 29。

---

4. 训练基础设施

我们的训练和评估代码库基于 JAX (Bradbury et al., 2018) 和 T5X (Roberts et al., 2022)，所有模型都在 TPU v4 Pod (Jouppi et al., 2020) 上进行训练。PaLM 540B 是在两个 TPU v4 Pod 上训练的，这两个 Pod 通过数据中心网络 (DCN) 连接，使用了模型并行和数据并行的组合 (Xu et al., 2021)。我们在每个 Pod 中使用 3072 个 TPU v4 芯片，连接到 768 个主机。该系统是迄今为止描述的最大 TPU 配置，使我们能够高效地将训练扩展到 6144 个芯片，而无需使用任何流水线并行 (Huang et al., 2019)。

先前关于类似规模端到端模型训练的报告使用了两种方法之一。LaMDA (Thoppilan et al., 2022) 和 GLaM (Du et al., 2021) 都是在单个 TPU 系统上训练的，没有利用流水线并行或 DCN。Megatron-Turing NLG 530B (Smith et al., 2022) 是在 2240 个 A100 GPU 上使用模型、数据和流水线并行的组合进行训练的，而 Gopher (Rae et al., 2021) 是在四个通过 DCN 连接的 TPU v3 Pod（每个 Pod 有 1024 个 TPU v3 芯片）上使用 Pod 间流水线进行训练的。

流水线通常与 DCN 一起使用 (Smith et al., 2022)，因为它具有较低的带宽要求，并提供了超出模型和数据并行所允许的最大高效规模之外的额外并行化。流水线通常将训练批次拆分为“微批次”，但它有重要的缺点。首先，它会产生流水线“气泡”的步长时间开销，其中许多设备在填充和清空前向和后向传递的流水线时完全处于空闲状态。其次，由于需要为小批次内的每个微批次从内存中重新加载权重，它需要更高的内存带宽。在某些情况下，它还涉及增加的软件复杂性。我们能够使用以下策略高效地将 PaLM 540B 的无流水线训练扩展到 6144 个芯片。

每个 TPU v4 Pod 包含模型参数的完整副本，每个权重张量分布在 3072 个芯片上，使用 12 路模型并行和 256 路完全分片数据并行（Xu et al. (2021) 中称为“2D finalized”的方法）。在前向传递期间，权重在数据并行轴上进行 all-gather，并且从每一层保存一个完全分片的激活张量。在后向传递期间，其余的激活被重新计算，因为这在更大的批量大小下会导致更高的训练吞吐量，与替代的重新计算选择相比。

我们使用 Pathways 系统 (Barham et al., 2022) 将训练扩展到单个 TPU v4 Pod 之外。PaLM 540B 利用 Pathways 的客户端-服务器架构在 Pod 级别实现双向数据并行。在这里，单个 Python 客户端将训练批次的一半分派给每个 Pod，每个 Pod 执行前向和后向计算，使用标准的 Pod 内数据和模型并行并行计算梯度。然后，Pod 与远程 Pod 传输梯度（在其批次的一半上计算），最后，每个 Pod 累积本地和远程梯度，并并行应用参数更新，以获得用于下一个时间步的位完全相同的参数。

图 2 显示了 Pathways 系统如何执行 Pod 级别的双向数据并行。单个 Python 客户端构建一个分片数据流程序（如图 2 左侧所示），该程序在每个包含 TPU Pod 的远程服务器上启动 JAX/XLA (XLA, 2019) 工作。该程序包含一个用于 Pod 内前向+后向计算（包括 Pod 内梯度归约）的组件 A，一个用于跨 Pod 梯度传输的传输子图，以及一个用于优化器更新（包括本地和远程梯度求和）的组件 B。Pathways 程序在每个 Pod 上执行组件 A，然后将输出梯度传输到另一个 Pod，最后在每个 Pod 上执行组件 B。Pathways 系统设计具有几个特性，使其能够将程序执行扩展到数千个加速器芯片——首先，它通过 Pod 内调度程序（如图 2 右侧所示）的异步组调度来掩盖从单个 Python 客户端向远程服务器分派 JAX/XLA 工作的延迟，其次，它通过分片数据流执行模型摊销了管理数据传输的成本（详情请参阅 Barham et al. (2022)）。

双向 Pod 级别数据并行的一个有趣方面是，在连接到两个 Pod 上总共 1536 个主机的 6144 个 TPU v4 芯片规模下，实现跨 Pod 梯度传输的高训练吞吐量的挑战。请注意，跨 Pod 梯度传输仅需要在两个 Pod 上的相应主机之间进行 1:1 传输，因为每个核心仅需要其模型分片参数的远程梯度。此外，两个 Pod 之间的主机通过 Google 数据中心网络 (Singh et al., 2015) 连接。由于传输直到每个核心完成梯度计算才开始（如图 2 所示），这导致了一个非常突发的工作负载，所有主机同时通过数据中心网络链路传输它们的梯度。特别是，每对主机在每个训练步中交换大约 1.3 GB 的梯度，这相当于所有主机上的总突发量为 81 Tbps。该工作负载的突发特性带来了挑战，我们通过仔细设计 Pathways 网络堆栈来解决这些挑战，以实现最佳的 DCN 链路利用率。例如，为了减轻拥塞的影响，梯度传输的数据被分解成更小的块，并通过一组多样化的 DCN 链路通过多个更小的流进行路由。通过这些优化，我们实现了训练吞吐量约为单个 Pod 训练吞吐量的 1.95 倍（相当于 97% 的完美弱扩展，因为我们相对于单个 Pod 将两个 Pod 的批量大小加倍）。与理论上的 2 倍吞吐量相比，性能差距源于后向传递与跨 Pod 梯度归约之间缺乏重叠。我们希望在未来的工作中解决这个问题。

4.1 训练效率

大多数先前报告的语言模型加速器效率数字使用我们称为硬件 FLOPs 利用率（HFU）的指标。这通常反映了在给定设备上观察到的 FLOPs 与其理论峰值 FLOPs 的比率估计。然而，硬件 FLOPs 利用率存在几个问题。首先，执行的硬件 FLOPs 数量是系统依赖和实现依赖的，编译器中的设计选择可能导致操作数量不同。重新计算（Rematerialization）是一种广泛用于权衡内存使用与计算的技术。为了高效地使用梯度下降计算大多数神经网络架构的后向传递，必须在内存中存储批次的许多中间激活。如果它们不能全部容纳，则可以重新计算一些前向传递操作（使一些激活能够被重新计算而不是存储）。这产生了一种权衡，即使用额外的硬件 FLOPs 可以节省内存，但训练系统的最终目标是实现每秒标记的高吞吐量（从而实现快速训练时间），而不是尽可能多地使用硬件 FLOPs。其次，测量观察到的硬件 FLOPs 取决于用于计数或跟踪它们的方法。观察到的硬件 FLOPs 已基于分析核算 (Narayanan et al., 2021b) 以及使用硬件性能计数器 (Xu et al., 2021) 进行了报告。

鉴于这些问题，我们认识到 HFU 不是 LLM 训练效率的一致且有意义的指标。我们提出了一种新的效率指标，该指标与实现无关，并允许更清晰地比较系统效率，称为模型 FLOPs 利用率（MFU）。这是观察到的吞吐量（每秒标记）相对于以峰值 FLOPs 运行的系统的理论最大吞吐量的比率。至关重要的是，“理论最大”吞吐量仅考虑计算前向+后向传递所需的必要操作，而不考虑重新计算。因此，MFU 允许在不同系统上的训练运行之间进行公平比较，因为分子只是观察到的每秒标记，而分母仅取决于模型架构和给定系统的已发布最大 FLOPs。我们在附录 B 中详细阐述了计算 MFU 的数学公式。

模型	参数数量（十亿）	加速器芯片	模型 FLOPs 利用率
GPT-3	175B	V100	21.3%
Gopher	280B	4096 TPU v3	32.5%
Megatron-Turing NLG	530B	2240 A100	30.2%
PaLM	540B	6144 TPU v4	46.2%

表 3：PaLM 和先前大型模型的模型 FLOPs 利用率。PaLM 实现了显著的高 MFU，因为在模型、编译器和并行策略方面进行了多项优化。PaLM 相应的硬件 FLOPs 利用率为 57.8%。计算详情见附录 B。

我们在表 3 中展示了 PaLM 540B 模型的模型 FLOPs 利用率，并将其与先前的大型模型进行了对比。MFU 在比较不同模型参数数量、架构和模型质量的背景下的模型和系统时非常有用。GPT-3 的 MFU 数字为 21.3%，基于 OpenAI 通过 Patterson et al. (2021) 报告的每个 GPU 24.6 非注意力模型 TFLOP/s，而 Gopher 的 MFU 数字为 32.5%，基于每秒 0.0152 步的训练速度。Megatron–Turing NLG 530B 的 MFU 数字为 29.7%（无自注意力）或 30.2%（有自注意力），基于每秒 65.43K 标记的训练吞吐量 (Smith et al., 2022)。

相比之下，PaLM 540B 在 2048 的批量大小下实现了每秒 238.3K 标记的平均训练吞吐量。PaLM 540B 的训练使用了重新计算，因为使用重新计算带来的更高可行批量大小实现了更高的训练吞吐量。PaLM 540B 的 MFU 为 45.7%（无自注意力）或 46.2%（有自注意力）。我们分析计算出的硬件 FLOPs 利用率（包括重新计算 FLOPs）为 57.8%。PaLM 实现了高加速器利用率，这得益于其并行策略和其他几个因素，包括 XLA TPU 编译器优化以及“并行层”的使用（见第 2 节）。我们相信 PaLM 代表了 LLM 训练效率的重大进步。

---

5. 训练设置

模型训练遵循大型 Transformer 语言模型的相当标准的设置。详情如下：

• 权重初始化 – 核权重（即嵌入和层归一化缩放之外的所有内容）使用“fan-in 方差缩放”进行初始化，即 $W \sim \mathcal{N}(0, 1/\sqrt{n_{in}})$，其中 $n_{in}$ 是核的输入维度。输入嵌入初始化为 $E \sim \mathcal{N}(0, 1)$，因为层归一化不应用于嵌入。由于输入和输出嵌入层是共享的，我们将 softmax 前的输出 logits 缩放为 $1/\sqrt{n}$，其中 $n$ 是嵌入大小。
• 优化器 – 模型使用 Adafactor 优化器 (Shazeer & Stern, 2018) 进行训练，没有进行因式分解。这实际上等同于具有“参数缩放”的 Adam (Kingma & Ba, 2014)，它通过参数矩阵的均方根来缩放学习率。由于权重初始化与 $1/\sqrt{n}$ 成正比，其效果类似于 Brown et al. (2020) 中 Adam 学习率的手动缩放。然而，参数缩放的好处是，在不同尺度下运行的参数矩阵（嵌入和层归一化缩放）不会以相同的速率缩放其学习率。
• 优化超参数 – 我们在前 10,000 步使用 $10^{-2}$ 的 Adafactor 学习率，然后以 $1/\sqrt{k}$ 的速率衰减，其中 $k$ 是步数。我们使用 $\beta_1 = 0.9$ 的动量进行训练。二阶矩插值值计算为 $\beta_2 = 1.0 - k^{-0.8}$，其中 $k$ 是步数。我们发现这在训练大型语言模型时比标准的 $\beta_2 = 0.99$ 更稳定，因为罕见的嵌入标记在较短的窗口内可能具有估计不佳的二阶矩。我们对所有模型使用值为 1.0 的全局范数梯度裁剪 (Pascanu et al. (2012))。我们在训练期间使用 $lr^{2.0}$ 的动态权重衰减，其中 $lr$ 是当前学习率。
• 损失函数 – 模型使用标准的语言建模损失函数进行训练，即所有标记的平均对数概率，不进行标签平滑。我们还使用了 $z_{loss} = 10^{-4} \cdot \log^2 Z$ 的辅助损失，以鼓励 softmax 归一化器 $\log(Z)$ 接近 0，我们发现这增加了训练的稳定性。
• 序列长度 – 所有模型都使用了 2048 的序列长度。输入示例被连接在一起，然后拆分为正好 2048 个标记的序列，因此没有填充标记，但示例可能会在中间被拆分。输入示例通过特殊的 [eod] 标记彼此区分。
• 批量大小 – 对于所有模型，我们在训练期间增加了批量大小。对于最大的模型，我们使用批量大小 512（1M 标记）直到 50k 步，然后将其加倍至 1024（2M 标记）直到 115k 步，最后再次加倍至 2048（4M 标记）直到在 255k 步完成训练。较小的模型遵循类似的计划。使用这种批量大小计划的原因有两个：（1）较小的批量大小在训练早期更具样本效率（即作为所见标记函数的损失更好），而较大的批量大小在训练后期由于更好的梯度估计而有益 (Smith et al., 2018; McCandlish et al., 2018)，以及 (2) 较大的批量大小导致更大的矩阵乘法维度，这提高了 TPU 效率。
• 位确定性 – 模型可以从任何检查点完全位可重现。换句话说，如果模型在单次运行中训练到 17,000 步，并且我们从 15,000 步的检查点重新启动，那么训练框架保证在从 15,000 到 17,000 的两次运行中产生相同的结果。这是通过两种方式实现的：（1）JAX+XLA+T5X 提供的位确定性建模框架，以及 (2) 确定性数据集管道，其中洗牌后的数据以随机访问格式写出，因此给定训练批次的内容仅是步数的函数 (Roberts et al., 2022)。
• Dropout – 模型在训练时没有使用 dropout，尽管在大多数情况下微调时使用了 0.1 的 dropout。

---

5.1 训练不稳定性

对于最大的模型，尽管启用了梯度裁剪，但我们在训练期间观察到损失出现了大约 20 次尖峰。这些尖峰发生在非常不规则的时间间隔，有时发生在训练后期，并且在训练较小的模型时没有观察到。由于训练最大模型的成本，我们无法确定缓解这些尖峰的原则性策略。

相反，我们发现了一种有效缓解该问题的简单策略：我们从尖峰开始前大约 100 步的检查点重新开始训练，并跳过了大约 200–500 个数据批次，这些批次涵盖了在尖峰之前和期间看到的批次。通过这种缓解措施，损失在同一点没有再次出现。我们不认为这些尖峰是由“坏数据”本身引起的，因为我们运行了几个消融实验，我们取了尖峰周围的数据批次，然后从不同的、更早的检查点开始在这些相同的数据批次上进行训练。在这些情况下，我们没有看到尖峰。这意味着尖峰仅由于特定数据批次与特定模型参数状态的组合而发生。在未来，我们计划研究在超大型语言模型中缓解损失尖峰的更具原则性的策略。

---

6. 评估

6.1 英语 NLP 任务

为了与先前的大型语言模型进行比较，我们在与 Du et al. (2021) 和 Brown et al. (2020) 相同的 29 个英语基准测试集上评估了 PaLM 模型。这些基准测试包括：

• 开放域闭卷问答任务：$^{5}$ TriviaQA (Joshi et al., 2017)、Natural Questions (Kwiatkowski et al., 2019)、Web Questions (Berant et al., 2013)
• 完形填空和补全任务：LAMBADA (Paperno et al., 2016)、HellaSwag (Zellers et al., 2019)、StoryCloze (Mostafazadeh et al., 2016)
• Winograd 风格任务：Winograd (Levesque et al., 2012)、WinoGrande (Sakaguchi et al., 2020)
• 常识推理：PIQA (Bisk et al., 2019)、ARC (Clark et al., 2018)、OpenBookQA (Mihaylov et al., 2018)
• 上下文阅读理解：DROP (Dua et al., 2019)、CoQA (Reddy et al., 2018)、QuAC (Choi et al., 2018)、SQuADv2 (Rajpurkar et al., 2018)、RACE (Lai et al., 2017)
• SuperGLUE (Wang et al., 2019a)
• 自然语言推理 (NLI)：Adversarial NLI (Nie et al., 2019)

表 4 包含了 PaLM 540B 的结果以及来自其他大型语言模型的先前最先进 (SOTA) 结果。在此表中，我们仅考虑来自预训练语言模型的单检查点结果。任何使用微调或多任务适配 (Wei et al. (2022a), Sanh et al. (2021)) 的模型均未包含在表中。

PaLM 540B 在 1-shot 设置下在 29 个任务中的 24 个上，以及在少样本设置下在 29 个任务中的 28 个上超越了先前的 SOTA。有趣的是，PaLM 540B 在一些阅读理解和 NLI 任务的少样本设置中，以超过 10 个点的优势超越了先前的 SOTA。虽然模型规模在实现这些结果方面发挥了重要作用，但 PaLM 540B 在所有基准测试中都优于类似规模的模型（Megatron-Turing NLG 530B）。这表明预训练数据集、训练策略以及训练期间观察到的标记数量在实现这些结果方面也发挥了重要作用。

$^{5}$ 在“闭卷”格式中，仅提供问题（以及示例的答案），而不提供上下文文档。Natural Questions 也有一个带有上下文文档的“开卷”格式，此处未使用。

---

表 5 列出了自然语言理解 (NLU) 和自然语言生成 (NLG) 任务的平均得分。PaLM 540B 在两个类别中的平均得分均提高了 5 分以上。如表所示，PaLM 模型的每个类别的平均得分也随着规模的增加而提高。有趣的是，PaLM 62B 在两个类别中都优于 GPT-3 175B。8B 和 62B 模型的所有任务结果可以在附录 H.1 中找到。

任务	0-shot (Prior SOTA)	0-shot (PaLM 540B)	1-shot (Prior SOTA)	1-shot (PaLM 540B)	Few-shot (Prior SOTA)	Few-shot (PaLM 540B)
TriviaQA (EM)	71.3$^a$	76.9	75.8$^a$	81.4	75.8$^a$ (1)	81.4 (1)
Natural Questions (EM)	24.7$^a$	21.2	26.3$^a$	29.3	32.5$^a$ (1)	39.6 (64)
Web Questions (EM)	19.0$^a$	10.6	25.3$^b$	22.6	41.1$^b$ (64)	43.5 (64)
Lambada (EM)	77.7$^f$	77.9	80.9$^a$	81.8	87.2$^c$ (15)	89.7 (8)
HellaSwag	80.8$^f$	83.4	80.2$^c$	83.6	82.4$^c$ (20)	83.8 (5)
StoryCloze	83.2$^b$	84.6	84.7$^b$	86.1	87.7$^b$ (70)	89.0 (5)
Winograd	88.3$^b$	90.1	89.7$^b$	87.5	88.6$^a$ (2)	89.4 (5)
Winogrande	74.9$^f$	81.1	73.7$^c$	83.7	79.2$^a$ (16)	85.1 (5)
Drop (F1)	57.3$^a$	69.4	57.8$^a$	70.8	58.6$^a$ (2)	70.8 (1)
CoQA (F1)	81.5$^b$	77.6	84.0$^b$	79.9	85.0$^b$ (5)	81.5 (5)
QuAC (F1)	41.5$^b$	45.2	43.4$^b$	47.7	44.3$^b$ (5)	47.7 (1)
SQuADv2 (F1)	71.1$^a$	80.8	71.8$^a$	82.9	71.8$^a$ (10)	83.3 (5)
SQuADv2 (EM)	64.7$^a$	75.5	66.5$^a$	78.7	67.0$^a$ (10)	79.6 (5)
RACE-m	64.0$^a$	68.1	65.6$^a$	69.3	66.9^a$$^\dagger (8)	72.1 (8)
RACE-h	47.9$^c$	49.1	48.7$^a$	52.1	49.3^a$$^\dagger (2)	54.6 (5)
PIQA	82.0$^c$	82.3	81.4$^a$	83.9	83.2$^c$ (5)	85.2 (5)
ARC-e	76.4$^e$	76.6	76.6$^a$	85.0	80.9$^e$ (10)	88.4 (5)
ARC-c	51.4$^b$	53.0	53.2$^b$	60.1	52.0$^a$ (3)	65.9 (5)
OpenbookQA	57.6$^b$	53.4	55.8$^b$	53.6	65.4$^b$ (100)	68.0 (32)
BoolQ	83.7$^f$	88.0	82.8$^a$	88.7	84.8$^c$ (32)	89.1 (8)
Copa	91.0$^b$	93.0	92.0$^a$	91.0	93.0$^a$ (16)	95.0 (5)
RTE	73.3$^e$	72.9	71.5$^a$	78.7	76.8 (5)	81.2 (5)
WiC	50.3$^a$	59.1	52.7$^a$	63.2	58.5$^c$ (32)	64.6 (5)
Multirc (F1a)	73.7$^a$	83.5	74.7$^a$	84.9	77.5$^a$ (4)	86.3 (5)
WSC	85.3$^a$	89.1	83.9$^a$	86.3	85.6$^a$ (2)	89.5 (5)
ReCoRD	90.3$^a$	92.9	90.3$^a$	92.8	90.6 (2)	92.9 (2)
CB	48.2$^a$	51.8	73.2$^a$	83.9	84.8$^a$ (8)	89.3 (5)
ANLI R1	39.2$^a$	48.4	42.4$^a$	52.6	44.3$^a$ (2)	56.9 (5)
ANLI R2	39.9$^e$	44.2	40.0$^a$	48.7	41.2$^a$ (10)	56.1 (5)
ANLI R3	41.3$^a$	45.7	40.8$^a$	52.3	44.7$^a$ (4)	51.2 (5)

表 4：PaLM 540B 模型在 29 个 NLP 基准测试中获得的结果。对于少样本结果，每个任务的 shot 数量在括号中提到。每个任务的拆分与 Du et al. (2021) 和 Brown et al. (2020) 中使用的相同。上标表示来自先前工作的结果：$^a$GLaM 62B/64E (Du et al., 2021), $^b$GPT-3 175B (Brown et al., 2020), $^c$Megatron-Turing NLG 530B (Smith et al., 2022), $^d$Gopher (Rae et al., 2021), $^e$LaMDA (Thoppilan et al., 2022) (结果来自 Wei et al. (2022a)), $^f$Chinchilla (Hoffmann et al., 2022)。$^\dagger$ Rae et al. (2021) 和 Hoffmann et al. (2022) 的工作在 RACE-m/h 上取得了更高的分数，但这些论文的作者指出，由于任务设置的差异，他们的分数不能与 GPT-3 和其他大型 LM 进行比较。我们遵循 GPT-3 风格的任务设置进行 RACE-m/h，并使用相同的设置与先前结果进行比较。