6144块TPU，5400亿参数，会改bug、解读笑话，谷歌刚刚用Pathways训练了一个大模型

会员服务 ·

6144块TPU，5400亿参数，会改bug、解读笑话，谷歌刚刚用Pathways训练了一个大模型

2022 年 4 月 5 日 机器之心

机器之心报道

机器之心编辑部

谷歌的下一代架构 Pathways 已经用来训练大模型了。

‍

随着规模的增加，模型在处理多个任务时的性能逐渐提高，而且还在不断解锁新的能力。

在探讨现有 AI 模型的局限时，谷歌人工智能主管 Jeff Dean 曾经说过，今天的人工智能系统总是从头开始学习新问题。最终，我们为数千个单独的任务开发了数千个模型。以这种方式学习每项新任务不仅需要更长的时间，而且还需要更多的数据，效率非常低。

在 Jeff Dean 看来，理想的发展模式应该是训练一个模型来做成千上万件事情。为了实现这一愿景，他所在的团队去年提出了一种名叫「Pathways」的通用 AI 架构。Jeff Dean 介绍说，Pathways 旨在用一个架构同时处理多项任务，并且拥有快速学习新任务、更好地理解世界的能力。前段时间，该团队终于公布了 Pathways 的论文。

论文写道，PATHWAYS 使用了一种新的异步分布式数据流设计。这种设计允许 PATHWAYS 采用单控制器模型，从而更容易表达复杂的新并行模式。实验结果表明，当在 2048 个 TPU 上运行 SPMD（single program multiple data）计算时，PATHWAYS 的性能（加速器利用率接近 100%）可以媲美 SOTA 系统。

谷歌 Pathways 系统架构概览。

有了强大的系统，接下来就是训练模型了。

在刚刚公布的论文——「PaLM: Scaling Language Modeling with Pathways」中，谷歌宣布，他们用 Pathways 系统训练了一个 5400 亿参数的大型语言模型——PaLM（Pathways Language Model）。

论文链接：https://storage.googleapis.com/pathways-language-model/PaLM-paper.pdf

这是一个只有解码器的密集 Transformer 模型。为了训练这个模型，谷歌动用了 6144 块 TPU，让 Pathways 在两个 Cloud TPU v4 Pods 上训练 PaLM。

强大的系统和算力投入带来了惊艳的结果。研究者在数百个语言理解和生成任务上评估了 PaLM，发现它在大多数任务上实现了 SOTA 少样本学习性能，可以出色地完成笑话解读、bug 修复、从表情符号中猜电影等语言、代码任务。

有网友感叹说，「终于知道谷歌开发 TPU 是用来干嘛的了。这种级别的自然语言理解，一旦被应用程序所利用，并变得足够高效并广泛使用。这将彻底改变谷歌所做的一切。拥有能够实现这一目标的芯片是非常有价值的，而那些陷入购买或租用英伟达芯片的公司将处于不利地位。」

PaLM 架构概览

PaLM 只包含解码器（每个时间步只能关注自身和过去的时间步），对一种标准的 Transformer 架构（(Vaswani et al., 2017)）做出了如下更改：

SwiGLU 激活

研究者使用 SwiGLU 激活 (Swish(xW) · xV) 用于 MLP 中间激活，因为研究表明，与标准 ReLU、GeLU 或 Swish 激活相比，SwiGLU 激活能显著提高质量。注意，在 MLP 中，这确实需要三个矩阵乘法，而不是两个，但 Shazeer (2020) 在计算等效实验中证明了质量的提升。

并行层

研究者在每个 Transformer 模块中使用「并行」方法，而不是标准的「串行」方法。具体来说，标准方法可以写成：

并行方法可以写成

由于 MLP 和注意力输入矩阵乘法可以融合，这里的并行方法可以让大规模训练速度提升 15%。消融实验显示，在 8B 的规模下，质量下降很小，但在 62B 规模下，质量没有下降，因此研究者推断，并行层的影响会在 540B 规模下达到 quality neutral。

多查询（Multi-Query）注意力

标准 Transformer 方法使用 k 个注意力头，其中每个时间步长的输入向量被线性投影成形状 [k，h] 的「查询」、「键」和「值」张量，其中 h 是注意力头大小。这里，键 / 值投影对于每个头是共享的，即「键」和「值」被投影到[1，h]，但是「查询」仍然被投影到形状[k，h]。此前有研究表明，这对模型质量和训练速度的影响呈中性，但在自回归解码时间上可以带来显著的成本节约。这是因为在自回归解码过程中，标准多头注意力在加速器硬件上的效率很低，因为键 / 值张量不在实例之间共享，并且一次只有单个 token 被解码。

RoPE 嵌入

研究者使用了 RoPE 嵌入而不是绝对或相对位置嵌入，因为 RoPE 嵌入已被证明在长序列长度上具有更好的性能。

共享输入 - 输出嵌入

研究者共享了输入和输出嵌入矩阵，这是在过去的工作中经常做的（但不是普遍的）。

No Biases

研究者在任何密集核或层 norm 中都没有使用 biases。他们发现，这可以增加大型模型的训练稳定性。

词汇表

研究者使用了具有 256k token 的 SentencePiece 词汇表，选择这个词汇表是为了支持训练语料库中的多种语言（没有过多的分词）。词汇表是从训练数据中生成的，研究者发现这提高了训练效率。

用 Pathways 训练一个 5400 亿参数的语言模型

PaLM 是谷歌首次大规模使用 Pathways 系统将训练扩展到 6144 块芯片，这是迄今为止用于训练的基于 TPU 的最大系统配置。研究者在 Pod 级别上跨两个 Cloud TPU v4 Pods 使用数据并行对训练进行扩展，同时在每个 Pod 中使用标准数据和模型并行。与以前的大多数 LLM 相比，这是一个显著的规模增长。

PaLM 实现了 57.8% 的硬件 FLOPs 利用率的训练效率，这是 LLM 在这个规模上实现的最高效率。为了达到这一水平，研究者将并行策略和 Transformer 块的重新设计结果相结合，这使得注意力层和前馈层并行计算成为可能，从而实现了 TPU 编译器优化带来的加速。

PaLM 使用英语和多语言数据集进行训练，包括高质量的 web 文档、书籍、维基百科、对话和 GitHub 代码。研究者还创建了一个「无损（lossless）」词汇表，它保留了所有空格（对于代码来说尤其重要），将词汇表之外的 Unicode 字符拆分成字节，并将数字拆分成单独的 token，每个 token 对应一个数字。

突破性的语言、推理和代码生成能力

PaLM 在许多非常困难的任务上显示出了突破性的能力，包括语言理解、生成、推理和代码等相关任务。

语言理解与生成

在 29 项基于英语的 NLP 任务上，PaLM 540B 的性能比之前的 SOTA 结果有所提高。

除了英语 NLP 任务外，PaLM 在包括翻译在内的多语言 NLP 基准测试中也表现出强大的性能，尽管它只有 22% 的训练语料库是非英语的。

研究者还在 Beyond the Imitation Game Benchmark (BIG-bench) 上探索了 PaLM 新的和未来功能，这是一个近期发布的套件，包含 150 多个新的语言建模任务。在这个过程中， PaLM 与 Gopher 和 Chinchilla 的性能进行了比较，涉及这些任务的 58 个公共子集。

有趣的是，研究者注意到，作为 scale 函数（function of scale）的 PaLM 的性能遵循与先前模型类似的对数线性表现，这表明 scale 的性能改进尚未趋于平稳。PaLM 540B 5-shot 的表现也优于解决相同任务的人类平均表现。

PaLM 在 58 个 BIG-bench 任务的子集上的 scaling behavior。

PaLM 在几个 BIG-bench 任务中展示了令人印象深刻的自然语言理解和生成能力。例如，该模型可以区分因果关系，理解上下文中的概念组合，甚至可以从表情符号中猜测电影。

PaLM 540B 在 BIG-bench 任务上的 1-shot 性能展示：标记因果关系、概念理解、从表情符号中猜测电影以及查找同义词和反事实。

推理

通过将模型 scale 与 chain-of-thought prompting 相结合，PaLM 在需要多步骤算术或常识推理的推理任务上展示出了突破性的能力。以往诸如 Gopher 这样的大型语言模型在提高性能方面从模型 scale 中获益较少。

小学数学问题示例中的标准 prompting 与 chain-of-thought prompting。Chain-of-thought prompting 将多步骤推理问题的 prompt 分解为中间步骤（黄色部分），类似于人类处理它的方式。

研究者在三个算术数据集和两个常识推理数据集上观察到了 PaLM 540B 在 chain-of-thought prompt 加持下的强大性能。例如，借助 8-shot prompt，PaLM 解决了 GSM8K 中 58% 的问题，这是一个包含数千个具有挑战性的小学水平数学问题的基准，超过了之前 GPT-3 175B 微调模型（训练集包含 7500 个问题，并与外部计算器和验证器相结合）获得的 55% 的最高分。

这个新的得分值得关注，因为它接近 60% 的 9 到 12 岁儿童解决问题的水平，这些儿童正是问题集的目标受众。研究者猜测，PaLM 词汇表中数字的独立编码有助于实现这些性能改进。

值得注意的是，PaLM 甚至可以为需要多步骤逻辑推理、世界认知和深度语言理解的复杂组合的场景生成明确的解释。例如，它可以为网络上搜不到的新笑话提供高质量的解释。

PaLM 用 two-shot prompts 解释了一个原创笑话。

代码生成

大型语言模型已被证明可以很好地推广到编码任务，比如在给定自然语言描述（文本到代码）的情况下编写代码，将代码从一种语言翻译成另一种语言，以及修复编译错误（代码到代码）。

PaLM 540B 在单个模型中显示了横跨编码任务和自然语言任务的强大性能，即使它在预训练数据集中只有 5% 的代码。具体而言，PaLM 540B 的 few-shot 性能十分显著，与经过微调的 Codex 12B 相当，同时使用的 Python 训练代码减少到了 50 分之一。这一结果印证了之前的发现，即较大的模型比较小的模型更高效，因为它们可以更好地从其他编程语言和自然语言数据中实现迁移学习。

PaLM 540B 微调模型在文本到代码任务（例如 GSM8K - Python 和 HumanEval）和代码到代码任务（例如 Transcoder）上的示例。

此外，通过在纯 Python 代码数据集上微调 PaLM ，模型进一步提高了性能，团队称之为 PaLM-Coder。如下图所示，PaLM-Coder 接到了一个名为 DeepFix 的示例代码修复任务，目标是修改最初损坏的 C 程序直到它们编译成功，PaLM-Coder 540B 展示了令人印象深刻的性能，实现了 82.1% 的编译率，优于之前 71.7% 的 SOTA 结果。这为修复软件开发过程中出现的更复杂的错误提供了机会。

DeepFix 代码修复任务的示例。经过微调的 PaLM-Coder 540B 将编译错误（左）修复为可编译的代码版本（右）。

从伦理方面考虑，最近的研究强调了受过网络文本训练的 LLM 相关的各种潜在风险。通过模型卡片和数据表等透明工件分析并记录这些潜在的不良风险是至关重要的，其中还包括有关预期用途和测试的信息。为此，谷歌的论文提供了数据表、模型卡片和 Responsible AI 基准测试结果，并报告了对数据集和模型输出的全面分析，以发现偏差和风险。虽然分析有助于概述模型的一些潜在风险，但特定领域和任务的分析对于真正校准、情境化和减轻可能的危害至关重要。进一步了解这些模型的风险和收益是正在进行的研究的主题，同时开发可扩展的解决方案，防止恶意使用语言模型。

参考链接：https://ai.googleblog.com/2022/04/pathways-language-model-palm-scaling-to.html