← 全部论文 / 2021

Scaling Language Models: Methods, Analysis and Insights from Training Gopher

2021 超大规模稠密 LLM 训练
下载 PDF

Scaling Language Models: Methods, Analysis and Insights from Training Gopher - 中文验证版

英文原文卡片：gopher_2021.md

状态：已翻译。

元数据

计算设置

论文明确使用 TPUv3 芯片，配合 JAX/Haiku 代码库和 JAX pmap 进行数据和模型并行。Gopher 是一个 280B 参数的自回归 Transformer，具有 80 层、128 头、模型维度 16,384、key/value 大小 128 以及 2048 token 的上下文窗口。它在 MassiveText 的 300B token 上训练，训练过程中 batch size 从 3M 增长到 6M token。7.1B 和 280B 模型使用 bfloat16 激活值和 bfloat16 参数；bfloat16 参数通过随机舍入更新。

附录给出了最具体的硬件规模。表 A27 列出了 280B 训练运行在 4096 个 TPUv3 芯片上。计算附录报告了 6.31E+08 训练 PFLOPs，并说明 Gopher 于 2020 年 11 月至 12 月在 Google 佐治亚数据中心训练了 920 小时。FLOP 核算包含实际成本，如重计算（rematerialization）增加 33% 计算量，加上 padding 和重复计算；不包括开发工作和抢占。

瓶颈

瓶颈是将一个密集的 280B 模型装入 TPUv3 内存和互连并使其移动。论文指出，半精度参数和单精度 Adam 状态占用 2.5 TiB，远超每个 TPUv3 核心可用的 16 GiB。这使得内存分区在有用训练开始之前就是强制性的。激活值也必须受到控制，这就是为什么重计算既表现为内存解决方案也表现为计算税。

这并非仅受参数量约束。一个通用混合精度 Adam 估计会将 280B 参数按每参数 16 字节计算为约 4.48 TB 的训练状态，而论文自己的核算报告了 2.5 TiB 用于半精度参数加单精度 Adam 状态。无论哪种方式，仅状态本身就碾压了一个 TPUv3 核心。对于推理，BF16 权重在 KV 缓存之前约为 560 GB；在 80 层、128 头、头维度 128 和 2048 token 上下文下，batch-1 的 KV 缓存大约再增加 10.7 GB，因此评估内存也是一个需要分片的服务问题。

在 4096 芯片时，通信和调度在时间分解中可见。对于 280B 模型在更高效的 6M token batch size 下，表 A27 将 51% 的加速器时间归因于线性层、8% 归因于注意力、3% 归因于优化器、7% 归因于模型并行、17% 归因于重计算、9% 归因于流水线、5% 归因于其他工作。因此瓶颈是内存容量、激活重计算、跨分片通信、流水线气泡和 batch size，而不仅仅是峰值 FLOPs。

方法适配

方法栈是对 TPUv3 pod 约束的直接回应。优化器状态分区、模型并行和重计算分区状态并减少激活值，使模型能装入内存。作者发现 TPUv3 上的数据和模型并行仅带来约 10% 的开销，因为跨芯片通信很快，因此在规模超出 1024 芯片 pod 之前无需流水线。对于 Gopher，他们在 TPU pod 内使用模型和数据并行，在 pod 之间使用流水线并行。

数值格式也是一种设备适配。bfloat16 减少内存并提高吞吐量，而随机舍入用于更新 bfloat16 参数。RMSNorm 和相对位置编码简化了架构，并允许在比训练时更长的序列上进行评估，无需重新训练上下文窗口。

数据是利用率的一个杠杆。MassiveText 包含 23.5 亿文档和约 10.5TB 来自网页、书籍、新闻、代码、C4 和 Wikipedia 的文本，但 Gopher 仅在 300B token 上训练，约为可用 token 的 12.8%。管线会过滤质量、去除重复、去重相似文档，并消除显著的测试集重叠。采样比例经过调优，MassiveWeb 占 48%，书籍 27%，C4 和新闻各 10%，GitHub 3%，Wikipedia 2%。

证据

训练证据异常详细。正文报告了 300B 训练 token、2048 上下文、Adam、batch size 从 3M 增长到 6M token、280B 模型使用 bfloat16、2.5 TiB 参数加 Adam 状态，以及 TPUv3 上约 10% 的数据/模型并行开销。附录补充了 4096 个 TPUv3 芯片、920 训练小时、6.31E+08 训练 PFLOPs，以及一个时间分解，其中仅重计算就消耗了 17% 的加速器时间。

基准证据展示了这种密集规模的回报。图 1 报告 Gopher 在 124 个可比较任务中的 100 个上改进了语言模型的现有最佳水平。在 RACE 阅读理解上，Gopher 在 RACE-high 上达到 71.6，对比 GPT-3 的 46.8 和 Megatron-Turing NLG 的 47.9；在 RACE-middle 上达到 75.1，对比 GPT-3 的 58.1。论文还指出 Gopher 在大多数涉及书籍和文章的语言建模任务上有改善，并将其与 MassiveText 较重的书籍采样联系起来。

计算附录也指出了推理的低效。它指出评估中对常见前缀的重复处理提高了推理成本，消除这种重复将根据评估任务将 FLOPs 降低 4 到 100 倍。在 280B 规模下的评估和提示本身就是不可忽视的计算工作负载。

历史影响

Gopher 既是一篇关于密集 LLM 系统的论文，也是一篇关于模型的论文。它记录了从"训练一个大的 Transformer"到"协调数千加速器芯片，同时分区优化器状态、重计算激活值和管理数据质量"的转变。其 280B 运行表明 TPU pod 规模的训练是可行的，但已经受到内存和通信的约束。

在历史上，它也奠定了 Chinchilla 批判的基础：一个在 300B token 上训练的 280B 模型产生了很强的结果，但后续研究认为该分配相对于计算最优扩展而言是训练不足的。Gopher 是前 Chinchilla 密集扩展范式的基线：在该范式中，优先选择最大参数量，而 token 数量相对较低。

局限

论文坦率地承认了残余局限。部分推理密集型和数学任务改进较少或退化；来源提到在 Ubuntu IRC 和 DM Mathematics 上表现不佳，并指出数字的 tokenizer 表示是可能的原因。安全和偏见问题仍然存在，包括毒性和子群体偏差分析。该模型训练和评估都很昂贵，而附录中关于重复前缀的讨论表明，简单的评估可能浪费大量推理计算。

之前的简短卡片指出确切的芯片数未被指定，但本地来源附录确实为 280B 训练分解指定了 4096 个 TPUv3 芯片。未指定的仍然是完整的开发计算、抢占/浪费开销以及所有服务细节。

链接

  • 计算范式：history/compute_regimes/hyperscale_dense_llm_training/README.md
  • 源 PDF 和抽取文本已在上方元数据中列出。
  • 队列状态：read_complete。