Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism 论文阅读

发表于 2025-04-05 更新于 2025-04-17 分类于 mlsys 阅读次数：本文字数： 2.7k 阅读时长 ≈ 10 分钟

NVIDIA 出品的大模型训练方法，2019 年挂在 ArXiV 上，到现在(2025 年)有 2k+ 论文引用。

#1. 引言

得益于计算能力和数据集大小的增长，NLP 领域取得了飞速的进展。

随着这些模型变得越来越大，它们超出了现代处理器的内存限制，需要额外的内存管理技术，例如 Activation Checkpointing。广泛使用的优化算法，如 Adam，需要每个参数存储额外的信息例如动量和其他优化器状态，这降低了可以训练的模型的大小。几种模型并行方法通过按比例划分模型参数来解决这个问题，使得权重和优化器状态不需要同时存在于处理器上。例如，GPipe 和 Mesh TensorFlow 提供了各种模型并行化框架。然而，这些方法需要重写模型，并且依赖于自定义的编译器和框架。

在这项工作中，作者实现了一个简单而有效的模型并行方法，使用层内模型并行。作者利用了 Transformer 架构的内存结构来实现简单的模型并行实现，可以高效地在 PyTorch 中进行训练，无需自定义 C++代码或者编译器。这种方法于 GPipe 等方法所倡导的基于管道的模型并行方法正交。

#2. 背景与挑战

#2.1. 神经语言模型预训练

预训练语言模型已经是 NLP 中不可或缺的组成部分。这些方法的进步需要高效的硬件、系统技术和框架。本文的工作旨在提供必要的工具，以便在这个趋势中迈出一步。

#2.2 Transformer 语言模型和多头注意力机制

由于其优越的准确性和计算效率，当前的 NLP 趋势是使用 Transformer 模型。原始的 Transformer 模型是一种机器翻译架构，分为编码器和解码器两个部分。然而，最近（2019 年）的工作，如 BERT 和 GPT-2 只使用了 Transformer 的编码器或解码器部分。

#2.3 深度学习中的数据并行和模型并行

有两种主要的方法来利用大量硬件加速器以扩展深度学习模型：数据并行（1990）和模型并行。数据并行方法中，每个 minibatch 被划分到多个 Worker 上。模型并行方法中，模型的内存使用和计算分布在多个 Worker 上。通过成比例增加 minibatch 大小和 Worker 数量的比例，可以观测到接近线性增长的训练数据吞吐量。然而，大的训练 batch 会引入优化过程的复杂性，可能导致准确性降低、或者收敛速度变慢，从而抵消了吞吐量增加的好处。 Parallel work 结合了数据并行和激活检查点方法来减少内存的需求。

然而，这些技术有一个共同的缺点：模型必须能完全放入 Worker 的内存中。随着语言模型越来越大和复杂，神经网络已经接近这个限制。解决这个问题的一个方法是使用参数共享(2019)来减少模型的内存占用，但是会限制模型的整体容量。本文的方法是使用模型并行来将模型分布在多个 GPU 上，这不仅缓解了内存压力，还增加了与 minibatch 大小无关的并行性。

在模型并行中，还有两个其他的范式：层间流水线并行，和更通用的分布式张量计算。在流水线模型并行中，在一个设备上先执行一组操作，然后输出会被传递到流水线上的下一个设备，在上面进行另一组操作。一些方法使用参数服务器和流水线并行相结合。然而会遇到不一致的问题，这种方法需要额外的处理逻辑，以及对优化器本身进行修改，而这些修改会降低效率或者影响准确性。

分布式张量计算是一种正交而且更加通用的方法，它将一个张量操作分布在多个设备上，以加速计算或者增加模型大小。 FlexFlow(Jia et al., 2018)提出了一种选择最优并行化策略的方法。最近，Mesh-TensorFlow(Shazeer et al., 2018)引入了一种语言来指定 TensorFlow 中的一般类分布张量计算。该语言由用户定义维度，使用适当的聚合原语编译生成的图。我们与 Mesh-TensorFlow 有类似的见解，但是我们没有实现新的框架和编译器，而是对现有的 PyTorch Transformer 实现进行了少量针对性的修改。我们的方法简单易懂，不需要任何新的编译器或者代码重写，通过插入一些简单的原语就可以完全实现，如下一节所述。

#3. 模型并行 Transformer

本文利用了 Transformer 的网络结构，通过添加几个同步原语来实现一个简单的模型并行。 Transformer 层由一个 Self Attention 后跟一个 MLP 组成，我们在两个块中分别引入了模型并行性。

首先介绍 MLP 块。MLP 块的第一部分是 GEMM，然后是 GeLU。

$Y=GeLU(XA)$

并行化 GEMM 的一种方法是按行拆分权重矩阵 $A$ ，按列拆分输入矩阵 $X$ ：

$X=[X_1, X_2] \\ A=\begin{bmatrix} A_1 \\ A_2 \end{bmatrix} \\ Y=GeLU(X_1 A_1 + X_2 A_2)$

因为 GeLU 是一个非线性函数，不能继续拆分成两个 GeLU 的和，这种方法需要在 GeLU 函数之前就插入同步点。

另一种选择是按列拆分 $A = [A_1, A_2]$ 。这种分区允许我们对每个分区的 GEMM 结果独立应用 GeLU 非线性：

$[Y_1, Y_2]=[GeLU(X A_1), GeLu(X A_2)]$

这种划分有利于消除同步点。因此，我们采用了这种方法。在分块进行第二次 GEMM 之后，进行 Reduce，然后进行 Dropout 层。这种方法将 MLP 块中的两个 GEMM 都分布在 GPU 上，并且只需要 Forward 和 Backward 各一次 All-Reduce。这两个操作彼此共轭，而且可以在 PyTorch 中使用几行代码实现。例如下面的代码实现：

class f(torch.autograd.Function):
    def forward(ctx, x):
        return x

    def backward(ctx, gradient):
        all_reduce(gradient)
        return gradient

对于 Self Attention，我们利用了 Multihead Attention 操作中的固有并行性，按照 K, Q 和 V 划分 GEMM 操作，这样每个 Attention 头对应的矩阵乘法都在一个 GPU 上本地执行。这使得我们能够将每个 Attention 头的参数和工作负载分布在多个 GPU 上，并且不需要任何即时通信就可以完成 Self Attention 的计算。输出 linear 层之后的 GEMM 按照行并行化，并直接采用并行注意力层的输出，无需 GPU 之间的通信。这种方法融合了两组 GEMM，去掉了中间一个同步点，并带来了更好的拓展性。这使得我们能够只使用 Forward/Backward 各两次 All-Reduce 的情况下，实现简单 Transformer 层中所有的 GEMM 计算。

Transformer 语言模型有一个 Hidden-size( $H$ ) x Vocabulary-size( $v$ ) 大小的输入 Embedding 层。由于现代语言模型的词表通常在万级别（例如，GPT-2 是 50,257 个），因此对输出 Embedding 层进行并行化是有好处的。然而，在 Transformer 语言模型中，输入和输出 Embedding 层共享权重，需要进行修改。我们按照单词维度并行化输入 Embedding 权重矩阵 $E_{H\times v} = [E_1, E_2]$ 。现在因为每个分片只包含 Embedding 表个一部分，在输入 Embedding 之后需要一次 All-Reduce。

对于输出 Embedding，一种办法是进行并行 $GEMM[Y_1, Y_2] = [X E_1, X E_2]$ 获取 logits，加上一次 All-gather 操作 $Y=all-gather([Y_1, Y_2])$ ，然后把结果发给 Cross-Entropy 损失函数。然而，All-gather 操作需要通信 $\text{batch_size} \times \text{seq_len} \times v$ 个元素，由于词表大小 $v$ 很大，总体通信也很大。为了减少通信大小，作者将并行 $GEMM[Y_1, Y_2]$ 的输出与 Cross-Entropy 聚合在一起，将维度降低至 $b \times s$ 。只通信标量的损失而不是 logits，极大地减少了通信量，提高了模型并行方法的效率。

我们的模型并发方法的主要特点是减少通信并保持 GPU 计算量。我们选择在多个 GPU 之间复制计算，而不是让一个 GPU 进行一部分计算然后广播结果。具体来说，我们在每个 GPU 上维护 LayerNorm 的参数，并在把参数输出到下一部分之前，在这些 Tensor 上进行 Dropout 和残差连接。为了优化模型，我们允许每个 Worker 优化自己的参数集合。由于所有的参数要么是 GPU 局部的，要么被拷贝过，因此不需要额外的通信去更新参数值。

简而言之，我们的方法实现简单，只需要在 Forward 和 Backward 的过程中增加一些额外的 All-Reduce 操作。不需要编辑器，而且与流水线模型并行方案正交。

#4. 设定

预训练语言理解模型是 NLP 和语言理解的核心任务。语言模型有几种方法。在本文中，我们关注 GPT-2，一种基于 Transformer 的自左向右的生成式语言模型，以及 BERT，一种基于语言模型掩码的双向 Transformer 模型。我们在下一部分解释这些模型的配置。

#4.1 训练数据集

包含 Wikipedia + CC-Stories + RealNews + OpenWebtext - WikiText103。 BERT 数据集额外包括 BookCorpus。过滤掉长度小于 128 个 token 的文档，使用 LSH 消除相似度大于 0.7 的重复内容。最终得到 174GB 的去重文本。

#4.2 训练优化，超参数

采用了带 Dynamic Loss Scaling 的 Mixed Precision 训练，以更好地利用 V100 的张量核心。权重初始化 $W~N(0, 0.02)$ 。在残差连接之前将权重乘以 $1/\sqrt{2N}$ ，其中 N 是 Transformer 层的数量。对于优化器，使用带 Weight Decay（ $\lambda=0.01$ ）的 Adam 优化器。使用了 1.0 的全局 Gradient Norm Clipping 来提升训练稳定性。 Dropout 使用 0.1。在每个 Transformer 之后使用 Activation Checkpointing。

对于 GPT-2 模型，Batch size 设定为 512，输入序列长度是 1024 个 subword，迭代 300k 次。学习率在前 3k 次迭代是 1.5e-4，后续采用单周期余弦衰减，在达到 1e-5 之后停止衰减。

对于 BERT 模型，使用了大小为 30522 的原始词典。使用 Sentence Order Prediction 替换了 Next Sentence Prediction 任务。使用整词 n-gram 掩码。 Batch size 设定为 1024，学习率为 1.0e-4，Warmup 10k 次，然后在剩下来的 2m 次迭代中线性衰减。

#5. 实验

TODO