← 全部论文 / 2017

Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer

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

Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer - 中文验证版

英文原文卡片：moe_2017.md

状态：已翻译。

元数据

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

计算设置

论文明确使用 GPU 集群。语言模型实验使用 16-32 块 Tesla K40 GPU 的集群。100B-word Google News 实验对大多数模型使用 32 块 Tesla K40 GPU，然后对最大两个模型使用 64 和 128 块 K40 GPU 以将参数装入内存。机器翻译实验在最多 64 块 GPU 上同步训练，每块 GPU 处理约 16,000 词的句对。翻译表中的 GNMT 比较使用 96 块 K80 GPU，多语言 GNMT 基线使用 96 块 K20s。

重要的设置细节是论文并非简单地将参数加到单个密集 GPU 模型上。密集循环层保持常规，而 MoE expert 是跨设备分布的模型并行分片。对 d 个设备上每个的 batch，MoE 层组合数据并行 batch，使每个 expert 看到的样本大约是直接每设备实现的 d 倍。

瓶颈

瓶颈是条件计算不自然匹配 GPU 硬件。论文称现代设备，尤其是 GPU，在算术上快但对分支不太友好，大 batch 是关键因为它们分摊参数传输和更新。直接的 MoE 将 batch 分割到许多 expert，因此对 n 个 expert 上的 top-k 路由，每个 expert 仅收到约 k×b/n 个样本。这恰好在模型试图扩展容量时产生微小矩阵乘和差的 GPU 占用。

网络带宽是另一瓶颈。expert 在设备上保持不动，因此通信主要包含将输入发送到选中的 expert 并返回输出。论文指出 GPU 集群的算术容量可以超过设备间聚合带宽数千倍。为保持计算受限而非网络受限，每个 expert 必须对每字节输入/输出执行足够多的算术。这解释了为何 expert 是具有数千 ReLU 单元的单隐藏层前馈网络，而非微小分支。

内存也约束 batch。更大 batch 使 expert 高效，但激活必须为反向传播存储。在 Google News 规模下，论文引入了激活重计算和分解二阶矩优化器状态，以在每块 GPU 上容纳最多约 1B 参数。

方法适配

稀疏门控 MoE 通过将每个样本路由到少量 expert（使用带噪声 top-k 门控）同时保持每个输入活跃 expert 数较少，将条件计算适配到 GPU 集群。在许多语言模型实验中 k=4；在一些翻译 MoE 中 k=2。结果是一个具有巨大参数容量的层，但每 token 计算路径更接近固定大小的密集层。

实现明确是混合并行的。非 MoE 层可以数据并行复制，而 expert 以模型并行的方式分片。来自许多设备的同步 batch 在 MoE 层合并，使 expert 子 batch 足够大以实现高效 GPU 矩阵乘。随着 expert 数量增长，设备数量可以增长，使每个 expert 的 batch size 和每设备内存保持可控。

门控损失既是硬件驱动的也是统计驱动的。如果门控坍缩到少数 expert，这些设备过载，expert batch 不平衡，大部分容量被浪费。重要性损失和负载均衡损失保持样本在 expert 间分散。在一些翻译实验中，作者进一步使用逐 batch mask，使训练期间每个 expert 收到相同 batch size，然后在推理时使用每 expert 阈值近似该行为。

证据

论文标题式声明是模型容量提高超过 1000x，而在现代 GPU 集群上计算效率仅小幅损失。实测效率证据混合但具体：低计算 MoE 模型运行在 0.74-0.90 TFLOPS/GPU，最高计算 MoE 达到 1.56 TFLOPS/GPU，最大 Google News 模型在 batch size 未随 GPU 数量成比例增加时降至 0.30 TFLOPS/GPU。最后一项结果有用，因为它暴露了 batch-size 瓶颈而非隐藏它。

在 100B-word Google News 数据集上，68.9B 参数 MoE-65536-h 模型在一个 epoch 后达到 28.9 测试 perplexity，约 0.72 观测 TFLOPS/GPU，对比 4xLSTM-512 基线的 47.0。137.7B 参数 MoE-131072-h 模型略恶化至 29.2 并降至 0.30 TFLOPS/GPU，表明该规模下稀疏度过大或 expert batch 过小。

在 WMT14 English-French 上，2048-expert MoE（8.7B 总参数）在 64 块 K40 上 6 天后达到 BLEU 40.56，对比 GNMT+RL 在 96 块 K80 上 6 天后 39.92。在 WMT14 English-German 上，相同规模在 64 块 K40 上 1 天后达到 BLEU 26.03，对比 GNMT 在 96 块 K80 上 24.91。在多语言翻译中，MoE 模型在 64 块 K40 上训练 12 天，有 8.7B 参数和 102M ops/timestep，在 12 个语言对中的 11 个上击败多语言 GNMT 基线。

历史影响

本文使条件计算在 GPU 集群上实用，并开启了稀疏容量分支：增加参数数量而不按比例增加每 token FLOPs。其历史贡献不仅是"使用 expert"，而是使 expert 可训练的具体计算配方：大共享 batch、expert 分片、负载均衡和通信感知的 expert 尺寸。

后来的 MoE 系统继承此结构。从本文到 GShard、Switch Transformer 和稀疏 LLM 的线索，是关于在内存和带宽约束下的容量。核心承诺是参数内存可以随设备扩展，而每 token 算术增长慢得多。

局限

极端稀疏时效率下降。最大 Google News 模型有更多参数，但 perplexity 更差、TFLOPS/GPU 远低于 68.9B 模型。论文将低效部分归因于未随 GPU 数量成比例增加 batch size，这强化了 MoE 缩放依赖足够大的 expert 子 batch。

方法还需要仔细的负载均衡和通信友好的 expert 形状。如果 expert 太小，网络流量占主导；如果路由坍缩，部分设备过载而其他闲置。推理也非免费：逐 batch 训练 mask 在服务时大训练 batch 不可用时需要近似。论文展示了稀疏容量，但非密集缩放的普适替代。

链接