← 全部论文 / 2020

GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding

2020 稀疏化与内存高效扩展
下载 PDF

GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding - 中文验证版

英文原文卡片：gshard_2020.md

状态：已翻译。

元数据

  • Slug： gshard_2020
  • 年份： 2020
  • 会议： arXiv
  • 作者： Dmitry Lepikhin et al.
  • 阅读状态： read complete
  • 计算范式： 稀疏化与内存高效扩展
  • 主要来源： PDF、抽取文本

计算设置

论文明确使用 TPU v3 集群：

  • 最多 2048 个 TPU v3 cores。
  • 2D TPU 集群结构。
  • Expert 数量绑定到设备数，例如 128、512 和 2048 cores。

主要报告的多语言模型是 MoE(2048E, 36L)，一个 600B 参数模型，在 2048 个 TPU v3 cores 上训练 4 天。表 3 报告每秒 0.72 步、4M-token batch、22.4 TPU v3 core-years 和平均 BLEU 44.3。论文还讨论了一个 1T 权重的 MoE(2048E, 60L) bfloat16 激活实验，但未纳入主要结果，因为可复现性需要仔细的手动诊断。

瓶颈

瓶颈是在不手写脆弱的模型并行代码的情况下扩展稀疏巨型模型。单设备内存对于数千亿参数来说太小，但朴素的模型并行让程序员协调跨设备通信，并可能使图大小或编译时间随分区数增长。论文明确指出模型在数量级上远超出单个加速器的内存容量，但也将手动分区视为软件瓶颈。

MoE 增加了自身的设备级瓶颈。Token 必须被路由到 experts，expert 缓冲区必须有界，dispatch/combine 步骤需要跨设备通信。AllToAll 是将 token 表示在组分片和 expert 分片布局之间移动的关键集合操作。注意力在短序列长度下仍受显存带宽约束，而密集前馈和投影层很好地映射到 TPU 矩阵单元。

方法适配

GShard 通过以下方式使 MoE Transformer 适配 TPU 集群：

  • 添加轻量级分片注释。
  • 使用 XLA SPMD 分区器。
  • 在交替前馈位置放置 top-2 MoE 层。
  • 对 token 分组以控制路由和 expert 容量。
  • 使用辅助负载均衡损失和随机第二 expert 路由。
  • 使用 AllToAll 进行 dispatch 和 combine。
  • 依赖编译器重物化。

分片 API 将模型描述与分区分开。用户可以用 replicate、split 或 shard 注释张量，XLA SPMD 分区器为所有设备生成单一程序。这避免了为每个分区生成独立的每设备图，并使编译可扩展到数千分区。MoE 模型主要在组/批次式分区和 expert 分区之间切换。

MoE 层使用 top-2 路由。一个可训练的门控网络为每个 token 选择最多两个 expert。Expert 容量由 token 数和 expert 数设置，溢出两个选定 expert 的 token 在该层被丢弃。为保持负载平衡，门控使用局部分组调度、辅助负载均衡损失和随机第二 expert 路由：如果第二 expert 的门控权重较小，该方法可以跳过它以节省 expert 容量。

实现是一种通信适配。Dispatch 和 combine 被表达为 einsum，然后编译器在维度之间插入 AllToAll 重分片。密集 Transformer FFN 和投影层仍是大矩阵乘法，以保持 TPU 利用率。当内存超过每设备限制时，编译器重物化在反向传播中重新计算激活。因此，模型花费额外计算以保持在内存内，并使用稀疏激活将参数计数与每 token 计算解耦。

证据

主要训练表给出标题结果。MoE(2048E, 36L) 使用 2048 cores，每秒 0.72 步运行，使用 4M-token batch，花费 22.4 TPU v3 core-years，训练 4.0 天，达到平均 BLEU 44.3，delta BLEU 13.5。512-expert、36 层模型使用 512 cores、1M-token batch、15.5 TPU core-years、11.0 天，BLEU 43.7。128-expert、36 层模型使用 128 cores、1M-token batch、6.1 TPU core-years、17.3 天，BLEU 39.0。

密集对比解释了为什么条件计算重要。表格列出密集 T(96L) 模型在 2048 cores 上 4M-token batch，约 235.5 TPU core-years，约 42 天，BLEU 36.9。引言还说明，具有 2.3B 参数的最佳密集单一 Transformer 基线达到低得多的质量，同时需要 235.5 TPU v3 core-years。因此，GShard 的 600B 稀疏模型大幅增加参数数量，同时相对于密集扩展路径降低了训练成本。

内存和运行时证据同等重要。随着设备和 expert 数量增加，每设备内存大致恒定，因为复制权重、分布式 MoE 权重和激活被 SPMD 布局分区。在固定层数下，权重内存和激活内存随 expert 数量增长保持恒定。步长时间次线性增长：论文报告从 128 扩展到 2048 设备/experts（expert 增加 16 倍）时执行时间仅增加 1.7 倍。

通信成为限制残余成本。在 128 expert 时，模型达到估计 roofline 性能的 70% 以上；在 2048 expert 时仍达到 48%。密集前馈层和 Transformer 投影达到 85% 以上峰值 FLOPs，而注意力达到 30% 以上，因为它更受显存带宽约束。MoE dispatch 和 combine 是 AllToAll 操作：当 expert 数量从 128 增加到 2048（16 倍），它们的执行时间增加约 3.75 倍，在 MoE 加 Transformer 时间中的份额从 16% 上升到 36%。通信微基准表明 AllReduce 在 TPU 上随设备数量大致恒定，而 AllToAll 大致随分区数的平方根缩放；从 16 到 2048 个分区，AllToAll 仅增加 9 倍，而分区数增加 128 倍。

历史影响

GShard 将 MoE 从手工构建的集群实验推进到编译器支持的 TPU 规模稀疏 Transformer。它使 600B 级条件计算模型可运行，并展示了模型容量、内存布局和通信集合操作可以通过轻量级注释和编译器 SPMD 分区来管理。

历史上，它桥接了原始的 sparsely gated MoE 思想和后来更简单的 expert 系统（如 Switch Transformer）。它还扩展了 Mesh TensorFlow 的教训：布局是方法的一部分，而非次要的实现细节。

局限

  • Expert 容量不能无限缩放，因为容量必须保持至少一个 token。
  • 1T bfloat16 激活模型仅通过仔细的手动诊断可训练，并存在数值稳定性问题，因此未纳入主要的可复现结果。
  • AllToAll 通信在运行时间中占比增长。
  • 系统依赖编译器和 mesh 假设。
  • 实验中将 expert 数量绑定到设备数以简化，尽管论文说这不是必需的要求。
  • 稀疏参数数量与每 token 密集计算不同；比较必须追踪激活计算、内存、通信和质量。

链接