Megatron-LM：借助优雅的 PyTorch 并行化，将语言模型扩展至数十亿参数

自然语言处理 (NLP) 领域已进入巨型模型的时代。从 GPT-2 到 BERT 及其后续发展，一个趋势十分明显: 模型越大，性能越好。这些庞大的 Transformer 模型不仅能够生成连贯的文章，还能回答复杂问题，并以前所未有的细腻程度解析语言。

但这种能力伴随着高昂的工程代价。这些模型拥有 数十亿——甚至越来越多地达到 数万亿——参数。如此庞大的网络如何能放入单个 GPU 的内存？

剧透: 它们放不下。

训练这些庞然大物需要先进的 并行化 技术，将模型实际拆分到多个 GPU 上运行。尽管谷歌的 GPipe 或 Mesh-TensorFlow 等框架提供了模型并行方案，但它们往往需要大量代码重写或定制编译器。

NVIDIA 研究人员推出了 Megatron-LM——一种利用 层内模型并行 训练超大 Transformer 模型的高效而简洁的方法。仅需在 PyTorch 中进行少量针对性修改，就能实现训练规模达 83 亿参数的 GPT-2，并提出了一个关键的架构调整，使 BERT 风格的模型也能平滑扩展。

在本文中，我们将探讨:

• 在 GPU 内存限制下训练超大模型的挑战
• Megatron-LM 优雅的层内模型并行方案
• 它如何实现近乎线性的扩展并达到业界领先的精度
• 一个微妙却至关重要的改动，解锁了大规模 BERT 训练

无论你是学生、研究人员，还是行业从业者，读完本文，你都将清楚理解支撑当今最大型语言模型的核心技术之一。

---

问题所在: GPU 的内存限制

现代 NLP 的核心是 Transformer 架构。你一定听说过其中的代表:

• BERT: 仅编码器的 Transformer，擅长理解上下文语义
• GPT-2: 仅解码器的 Transformer，擅长文本生成

图 2. Megatron-LM 中使用的 Transformer 架构。每一层交替包含自注意力模块和 MLP 模块，并由残差连接和层归一化包裹。

Transformer 是由多个相同结构的层堆叠而成，每一层包含:

1. 自注意力 (Self-Attention) 模块
2. 多层感知机 (MLP) 模块

增加层数和宽度 (隐藏单元数) 会让参数量急剧膨胀。

最常见的扩展方法是 数据并行: 将整个模型复制到多个 GPU 上，将批次数据拆分分配给这些 GPU，并对梯度求平均。这方法有效但有上限: 模型必须能放入单个 GPU 内存，参数量达到数十亿时就行不通了。

模型并行通过将同一个模型拆分到多个 GPU 上来规避这一限制。主要有两类:

1. 层间 (流水线) 并行: 不同层分布在不同 GPU 上。流水线方案如 GPipe 性能高，但容易出现“流水线气泡”，即 GPU 空闲等待时间。
2. 层内并行: 将每层内部的计算拆分到多个 GPU 上。这是 Megatron-LM 的核心——也是它的优势所在。

Megatron-LM 的关键洞见是: 利用 Transformer 组件的数学结构，实现简单而高效的层内并行。

---

核心方法: 切分 Transformer 层

每个 Transformer 层都有一个 MLP 模块和一个自注意力模块。Megatron-LM 对两者都做了高效的并行化处理。

并行化 MLP 模块

MLP 模块的计算公式为:

\[ Y = \operatorname{GeLU}(XA) \tag{1} \]

接着:

\[ Z = YB \]

其中 \(X\) 为输入，\(A\)、\(B\) 为权重矩阵，GeLU 是非线性激活函数。

如果对 \(A\) 按行切分到不同 GPU，需要在执行 GeLU 之前对结果求和——这会在模块中引入一次昂贵的中间同步。Megatron-LM 选择将 \(A\) 按列切分:

\[ [Y_1, Y_2] = [\operatorname{GeLU}(XA_1), \operatorname{GeLU}(XA_2)] \tag{3} \]

每个 GPU 独立计算自己的切片，此时无需通信。

第二个 GEMM 中，\(B\) 按行切分。每个 GPU 将其本地 \(Y_i\) 与 \(B_i\) 相乘，再在乘法之后通过一次 all-reduce 合并结果。

图 3. Megatron-LM 的并行化策略: MLP 先按列再按行切分；自注意力将注意力头分到不同 GPU 上。

该方法使 MLP 前向传播和反向传播各仅需 一次 all-reduce。在 PyTorch 中通过两个自定义 autograd 函数实现:

• g: 前向执行 all-reduce，反向为恒等操作
• f: 前向为恒等操作，反向执行 all-reduce

---

并行化自注意力模块

多头注意力天生适合并行: 每个头独立完成 Q、K、V 投影。Megatron-LM 将注意力头的子集分配给不同 GPU。投影矩阵按列切分到各 GPU，注意力分数按头在本地计算。

自注意力后的最终线性投影按行切分，并进行一次 all-reduce。

---

每层的通信开销

MLP 和自注意力并行化后，每个 Transformer 层总共只需 四次通信:

• 前向传播中 2 次 all-reduce
• 反向传播中 2 次 all-reduce

图 4. 模型并行的 Transformer 层中的通信操作。

嵌入层 (规模可能非常大) 也采用类似的按列切分方法进行并行化。

---

与数据并行结合

层内模型并行可与数据并行无缝配合。例如，使用 512 个 GPU:

• 分成 64 个“模型并行组”，每组 8 个 GPU
• 每组持有模型的一份切片
• 64 个组处理不同的小批量数据 —— 即标准数据并行

图 8. 混合方案: 纵向分组运行模型并行切片，横向分组运行数据并行副本。

这种 混合并行 能扩展到数百甚至上千 GPU。

---

实验与结果: 效率的证明

它能扩展吗？答案是肯定的——而且表现惊人。

扩展性能

团队测试了从 12 亿参数 (可放入单个 GPU) 到 83 亿参数 (需要 8 个 GPU 模型并行) 的 GPT-2 配置。

表 1. 扩展性研究的参数配置。更大模型需更多模型并行 GPU。

在 512 个 GPU 上采用混合并行，持续达成 15.1 PetaFLOPs。

图 1. GPU 增加时的近线性弱扩展性。绿色: 模型+数据并行；蓝色: 模型并行。

单 GPU 基线达到 39 TFLOPs (理论峰值的 30%) 。扩展到 512 GPU 时达到 76% 效率——通信损耗极低。

图 5. 模型并行和混合并行的弱扩展效率。

---

扩展准确率: GPT-2

团队训练了以下规模的 GPT-2 模型:

• 3.55 亿参数
• 25 亿参数
• 83 亿参数

表 2. GPT-2 模型配置。

更大的模型收敛更快，验证困惑度更低:

图 6. 验证困惑度与迭代次数的关系: 模型越大，学习越快。

零样本基准测试显示其达到业界领先水平:

表 3. SOTA 结果: 83 亿参数模型在 WikiText103 和 LAMBADA 上均优于此前模型。

---

BERT 的突破: 一个微妙但关键的改动

将 BERT 扩展到超过 3.36 亿参数一直有问题——训练不稳定且性能下降阻碍了进步。

Megatron-LM 找到原因:** 层归一化的位置**。

标准 BERT 中，层归一化在残差连接之后 (图 7a) 。若将层归一化移到每个模块的开头，即自注意力和 MLP 之前 (图 7b) ，训练稳定性显著提升。

图 7. 左: 原始与修改后的 BERT 模块结构。右: 修改后较大模型的损失曲线显著改善。

借助这一改动，他们将 BERT 训练规模提升到 39 亿参数:

表 4. Megatron-BERT 配置。

在各类下游任务上的性能也持续提升:

表 5. 更大的 Megatron-BERT 模型优于小规模模型；39 亿参数模型在 RACE 上达成 SOTA。

---

结论与启示

Megatron-LM 论文是高效大规模模型训练的里程碑。作者们通过少量有针对性的修改，在 PyTorch 中构建了一个可扩展、优雅的层内模型并行框架。

核心启示:

1. 层内模型并行很强大: 先按列后按行切分 GEMM，可最大限度减少通信，实现超大规模扩展
2. 简单性至关重要: 无需新编译器或大量代码重写，仅依靠巧妙工程设计即可
3. 架构调整可以解锁规模: 对 BERT 层归一化的小调整即可解决重大扩展难题
4. 开源推动进步: Megatron-LM 的发布让研究者可以自行训练巨型模型

自此，Megatron-LM 技术为更大模型 (如 5300 亿参数的 Megatron-Turing NLG) 奠定了基础，影响了众多并行训练系统。这证明了深入理解架构可以带来系统性能的重大突破。

---