← 全部论文 / 2019

Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism

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

Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism - 中文验证版

英文原文卡片：megatron_lm_2019.md

状态：已翻译。

元数据

  • Slug： megatron_lm_2019
  • 年份： 2019
  • 会议： arXiv
  • 作者： Mohammad Shoeybi et al.
  • 阅读状态： read complete
  • 计算范式： 超大规模密集 LLM 训练
  • 主要来源： PDF、抽取文本

计算设置

论文明确给出 NVIDIA 硬件：

  • 单 GPU 基线：DGX-2H 服务器中的 NVIDIA V100 32GB GPU。
  • 完整运行最多使用 32 台 DGX-2H 服务器，总计 512 块 Tesla V100 SXM3 32GB GPUs。

基础设施细节很重要：每台 DGX-2H 服务器通过 NVSwitch 在 GPU 之间提供 300 GB/sec 带宽，服务器之间通过每台 8 个 InfiniBand 适配器以 100 GB/sec 互联带宽连接。最大的类 GPT-2 配置使用 8 路张量模型并行和 64 路数据并行，共 512 块 GPUs。论文的强单 GPU 基线是一个在一块 V100 32GB GPU 上的 1.2B 参数模型，训练期间维持 39 TFLOP/s。

瓶颈

瓶颈是密集 Transformer 的内存和通信。一旦计入参数、优化器状态、激活和临时 buffer，数十亿参数的语言模型就会超过一块 V100 的内存。仅靠数据并行不会减少模型内存，而简单地增加 batch size 以提高吞吐量可能损害收敛。流水线并行可以减轻内存压力，但会引入调度气泡和实现复杂性。

训练内存的下限远大于原始参数文件。混合精度 Adam 在计入激活之前通常每个参数需要约 16 字节：半精度参数和梯度加上 FP32 主参数和 Adam 动量。按照这一计算，8.3B 的类 GPT 模型意味着约 133 GB 的模型状态，超过四块 32GB V100 的容量，还未计入激活或临时 buffer。即使是 1.2B 的单 GPU 基线也消耗约 19 GB 的模型状态，在一块 32GB V100 上留给序列激活、dropout、checkpointing buffer 和框架开销的空间有限。

Megatron-LM 关注一个更窄的硬件事实：在 DGX-2H 内部，V100 之间以非常高的带宽连接，因此层内张量并行可以拆分大型矩阵乘法，同时保持 GPU 计算受限。该方法在可能时避免将细粒度模型并行通信推过慢速链路，并将张量并行组与数据并行结合以实现全集群规模。

方法适配

Megatron-LM 通过以下方式将 Transformer 适配到 V100 集群：

  • 跨 GPU 拆分 Transformer MLP 矩阵乘法。
  • 跨 GPU 拆分注意力头。
  • 将层内张量模型并行与跨副本的数据并行结合。
  • 使用激活 checkpointing 减少内存。
  • 通过使用层内并行避免流水线气泡。

MLP 拆分围绕 GEMM 和 GeLU 设计。第一个线性层按列拆分，使得每个 GPU 计算其分区并局部应用 GeLU，在非线性之前避免了一次同步点。第二个线性层按行拆分，然后在 GPU 之间进行归约（reduction）。论文说这在 MLP block 的正向传播中需要一次 all-reduce，反向传播中需要一次。

自注意力按头拆分。key、query 和 value GEMM 被分区，使得每个注意力头的矩阵乘法局限于一个 GPU。然后输出投影被并行化以直接接收并行注意力输出。对于一个 transformer 层，论文将结果总结为正向路径两次 all-reduce，反向路径两次 all-reduce。这就是核心的计算适配：添加少量大规模结构化集合通信，以保持 tensor core 执行大规模 GEMM。

embedding 和 loss 也是通信感知的。由于 GPT-2 词汇量是数万个 tokens，all-gather 完整 logits 会传输 batch × sequence × vocabulary 的元素。Megatron 将 embedding 按词汇并行化，并将并行输出 GEMM 与交叉熵融合，使通信仅限于标量损失而非完整 logits。论文还将词汇填充到 51,200，使得在最多 8 路模型并行下，每个 GPU 的词汇大小是 128 的倍数，这更好地适配高效的 GEMM 形状。

内存通过混合精度、动态损失缩放、Adam、梯度裁剪、dropout 以及每个 transformer 层后的激活 checkpointing 来管理。这些不是附带的训练技巧；它们让密集 Transformer block 能够装入 V100 内存，同时保持 tensor-core 吞吐量高。

证据

缩放证据是明确的。论文训练了高达 8.3B 参数的类 GPT-2 模型，使用 512 块 GPUs，其中 8.3B 运行使用 72 层、hidden size 3072、24 个注意力头和 512 块总 GPUs。报告的 8.3B 模型的 epoch 时间约为 2.10 天。完整应用维持 15.1 PFLOP/s，相对于强单 GPU 基线达到约 74-76% 的缩放效率，具体取决于报告摘要/图表。

单 GPU 基线很重要，因为它并不弱：1.2B 模型在一块 V100 32GB GPU 上维持 39 TFLOP/s，约为理论峰值的 30%。以 8 路模型并行缩放 8.3B 模型达到线性缩放的 77%，而以模型加数据并行在 512 块 GPUs 上仍然达到线性缩放的 74%。模型加数据并行情况下的固定全局 batch size 为 512，对应于 64 路数据并行。

语言模型结果表明硬件缩放不仅仅是合成的吞吐量。8.3B 的类 GPT-2 模型达到 WikiText-103 perplexity 10.81 和 LAMBADA 完形填空准确率 66.51%，相对于较小的 355M 和 2.5B 模型有所提升。论文还训练了高达 3.9B 参数的类 BERT 模型，并在摘要中报告 RACE 准确率 90.9%。训练语料为 174GB 去重文本，由 Wikipedia、CC-Stories、RealNews 和 OpenWebText 构建，并对评估数据集进行了泄漏过滤。

历史影响

Megatron-LM 使密集 Transformer 缩放成为 V100 集群上实际的张量分区问题。它是超大规模 LLM 训练栈的关键前身，其中张量并行、激活 checkpointing、混合精度、优化器内存和互联拓扑都成为模型配方的组成部分。

其历史影响也是方法论层面的：论文表明大型 Transformer 层具有足够的内部结构，仅需少量集合通信即可拆分，无需新的编译器或对框架的完全重写。这使张量并行成为后来的系统将其与流水线并行、ZeRO 式优化器分区和数据并行结合的标准维度。

局限

  • 未来更大的模型需要超过单个 DGX-2H 机箱可用的内存。
  • 论文指向了混合层内、层间和节点间并行。
  • 它本身不解决 optimizer-state 冗余；ZeRO 单独处理该问题。
  • 该方法依赖非常高的节点内带宽；当细粒度模型并行通信穿越较慢的节点间链路时效率下降。
  • 通信虽被最小化但未被消除：每个模型并行的 transformer 层仍然执行 all-reduce。
  • 密集缩放仍然昂贵，因为每个参数参与每个 token。

链接