《Megatron-LM：用模型并行训练数十亿参数的语言模型》

大模型训练的内存瓶颈

近期的工作表明，训练大型 Transformer 模型能够推动自然语言处理的最先进水平，但超大模型会因内存限制而难以训练^[1]。训练时，模型的参数、参数的梯度以及优化器状态都要占用存储空间，而单块加速器的容量是固定的。一旦模型超出这个容量，就必须把一部分工作分散到多个设备上。

这篇论文的作者是 NVIDIA 的 Mohammad Shoeybi、Mostofa Patwary、Raul Puri、Patrick LeGresley、Jared Casper 和 Bryan Catanzaro，2019 年 9 月 17 日以 1909.08053 的编号提交到 arXiv^[1]。

把一层拆到多张卡上

作者提出一种简单、高效的层内模型并行（intra-layer model parallelism）方法，使训练参数量达到数十亿的 Transformer 模型成为可能^[1]。所谓模型并行，笼统地说就是把同一个模型分散到多台设备上，而不是在每台设备上各放一份副本：不同的计算片段放在不同的芯片上，设备之间交换彼此需要的中间结果。层内并行意味着切分发生在层的内部，一次大的矩阵乘法本身就被分摊到多张卡上。

作者强调的是这套方法对配套工具的要求很低：它不需要新的编译器，也不需要修改库，与流水线模型并行（pipeline model parallelism）相互正交、可以互补，只要在原生 PyTorch 中插入少量通信操作就能实现^[1]。换句话说，改动落在模型代码里而不是工具链里，并且它可以和其他切分模型的方式叠加使用，而不是取代它们。

规模与吞吐

作者用 512 块 GPU 收敛了参数量最高达 83 亿的 Transformer 模型，以此说明方法的可行性^[1]。论文报告，整个应用可以维持 15.1 PetaFLOPs 的算力；与一个强劲的单 GPU 基线相比——该基线维持 39 TeraFLOPs，相当于峰值算力的 30%——扩展效率为 76%^[1]。这里的扩展效率衡量的是：同样的工作分散到多台设备后，单设备的吞吐量还剩下多少。注意对照的是那个强劲的单 GPU 基线，而不是硬件的理论峰值。

两个大模型与基准成绩

为了说明大语言模型可以进一步推动最先进水平，作者训练了一个类似 GPT-2 的 83 亿参数模型和一个类似 BERT 的 39 亿参数模型^[1]。论文还报告了一项结构性发现：在类 BERT 模型中，仔细斟酌层归一化（layer normalization）放在什么位置，对于模型规模增大时取得更高性能至关重要^[1]。层归一化是 Transformer 模块内部对激活值重新缩放的一步；随着模型变大，它相对于其他子层所处的位置会变得越来越重要。

论文给出的数字是与此前最先进结果对照的。类似 GPT-2 的模型在 WikiText103 上的困惑度（perplexity）为 10.8，此前的最先进结果为 15.8；在 LAMBADA 上的准确率为 66.5%，此前为 63.2%；类似 BERT 的模型在 RACE 上达到 90.9%，此前为 89.4%^[1]。困惑度衡量的是语言模型对未见文本感到「意外」的程度，数值越低越好；另外两项是准确率，越高越好。

为什么重要

这篇论文给出了两个层面的结论。工程层面：参数量达到数十亿的模型可以用层内切分来训练，既不需要新编译器也不需要改库，只需在原生 PyTorch 中插入少量通信操作，并且能与流水线模型并行共存^[1]。实证层面：在留意层归一化位置这类架构细节的前提下把规模做大，作者所测的基准上，无论 GPT-2 式还是 BERT 式的设计，结果都优于此前的最先进水平^[1]。两点合起来，描述的是一条从「单卡放得下」走到「单卡放不下」的可行路径。