← 全部论文 / 2021

GLaM: Efficient Scaling of Language Models with Mixture-of-Experts

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

GLaM: Efficient Scaling of Language Models with Mixture-of-Experts - 中文验证版

英文原文卡片：glam_2021.md

状态：已翻译。

元数据

计算设置

论文给出了一个异常具体的训练设置。最大模型 GLaM 64B/64E 在 1024 个 Cloud TPU-v4 芯片上训练，大模型权重和计算通过 GSPMD 实现的 2D 分片算法分区。训练使用 float32 模型权重和 bfloat16 激活值，因此加速器适配策略不仅是"更多芯片"；它同时是一个精度和编译器布局的方案。优化器为 Adafactor，序列长度为 1024 token，并且样本经打包后，使一个 batch 可包含最多 100 万个 token。架构表中最大的行具有 1.2T 总参数、每 token 96.6B 激活参数、64 个 Transformer 层和 64 个专家。

能耗数据同样是陈述而非推断得出的。完整的 600B token 的 GLaM 64B/64E 训练运行据报告为 456 MWh，对比 GPT-3 的 1287 MWh。为达到与 GPT-3 相似且略优的分数，论文报告在 1024 个 TPU-v4 芯片上运行 574 小时处理 280B token，消耗 213 MWh。未指定 GPU 时代的推理硬件；具体的设备声明为 TPU-v4 训练加上 GSPMD 分区。

瓶颈

瓶颈是密集扩展中内存容量、训练计算和服务成本的耦合。像 GPT-3 这样的密集 175B 参数模型必须为每个 token 访问相同的参数集，因此增加总容量会直接增加每个 token 的 FLOPs。GLaM 针对的场景是：质量受益于大得多的参数储库，但延迟和能耗由每个 token 的活动子网络决定。

第二个瓶颈是分布式系统利用率。稀疏 MoE 将困难转移到 routing、专家负载均衡以及跨 TPU 设备网格的专家权重分区。如果 token 聚集在相同的专家上，设备会过载，batch 利用率下降；如果专家过大，专家权重本身也必须分片。因此 GLaM 是一篇计算结构论文：它探讨如何使用万亿参数的内存，且每次预测无需付出万亿参数 FLOPs 的代价。

方法适配

GLaM 通过将每隔一个前馈层替换为稀疏激活的 MoE 层来适配仅解码器的 Transformer。每个 token 路由到 64 个专家中的 top-2，仅这两个专家的前馈网络被评估。这保留了一个密集的 Transformer 主干，同时使昂贵的前馈路径以 token 为条件。结果是模型具有 1.2T 存储参数，但每 token 约 8% 处于活跃状态。

该方法根据来源中所述的内存和并行约束进行了调优。Top-2 路由是一个明确的质量/效率权衡：选择更多专家会提高每 token FLOPs 并使网络更密集，而选择更少专家会减少计算但可能损失建模质量。来自 GShard 的辅助 MoE 损失促使门控函数将 token 均匀分配到各专家，这是对设备不均衡的直接防御。2D 分片方法在 batch、模型和专家维度上分区权重和计算，使大专家能够跨多个 TPU 核心而无需复制所有数据或计算。bfloat16 激活值降低了激活带宽和内存压力，而 float32 权重使大型稀疏模型在数值上保持保守。

证据

核心证据是 GPT-3 对比表。GLaM 每 token 使用 180 GFLOPs，而 GPT-3 为 350 GFLOPs，减少 48.6%，同时论文报告了在该表所示基准套件上更好的平均零样本准确率。同一张表报告了 GLaM 训练能耗为 456 MWh，对比 GPT-3 的 1287 MWh。架构表给出了这些数字背后的机制：1.2T 总参数但仅 96.6B 活跃参数每 token。

扩展实验支持了计算论证，而不仅仅是最终分数。对于固定的每预测有效 FLOPs，增加更多专家通常会改善性能，这表明稀疏内存容量可以在无需等比例增加活跃 FLOPs 的情况下购买质量。论文还比较了在相同数据上训练的密集和 MoE 版 GLaM 变体，显示稀疏模型在类似有效计算量下实现了更强的 NLG 和 NLU 曲线。作者强调 280B token 的 GLaM 运行可以在 213 MWh 下达到或略超 GPT-3 级别的分数，这将基准质量与一个实测的设备能耗声明联系起来。

历史影响

GLaM 是介于 GShard/Switch 式 MoE 研究以及后来生产级 MoE 大语言模型之间的一个主要的稀疏扩展节点。它使稀疏扩展命题有了一个强有力的具体版本：总参数量可以增长到万亿范围，而活跃推理计算保持在接近一个小得多的密集模型的水平。它还帮助将评估对话从单纯的参数数量转向活跃参数、每 token FLOPs 和训练能耗。

在历史上，本卡片属于稀疏化与内存高效扩展范式，因为其核心主张不是 TPU 集群更大，而是架构以不同的方式使用集群：密集扩展同时消耗内存和计算，而 GLaM 通过稀疏激活和编译器支持的分片将二者分离。

局限

局限主要是系统层面的。服务 GLaM 仍然需要存储和移动一个 1.2T 参数的模型，即使每个 token 仅激活 96.6B 参数。MoE 引入了 routing 决策、辅助均衡损失、专家容量边界情况以及密集 Transformer 所避免的跨设备通信路径。论文中最强的硬件证据是针对 TPU-v4 加上 GSPMD 的；对其他互连或服务栈的可移植性未被本卡片来源确立。

与 GPT-3 的对比并非同一实验室重跑；它依赖于已发表的 GPT-3 和 Patterson 能耗估计。尽管如此，在论文证据范围内，稀疏激活将加速器内存容量转化为质量，而无需为每个 token 激活所有参数。

链接

  • 计算范式：history/compute_regimes/sparse_memory_efficient_scaling/README.md
  • 源 PDF 和抽取文本已在上方元数据中列出。
  • 队列状态：read_complete。
  • 方法索引：MoE、scaling laws
  • Ledger 更新：计算瓶颈