模型并行 | 大规模语言模型架构 Megatron

随着AI模型的规模越来越大，分布式训练技术越来越被广泛使用。现行的分布式训练方法主要包含两个部分：数据并行（Data Parallel）和模型并行（Model Parallel）。数据并行是将模型完整拷贝到多张显卡中，对批次数据进行并行计算，适合规模小而数据多的训练场景；而模型并行适合超大规模参数的模型训练，将模型不同的部分分别加载到不同的显卡中，依次计算得出结果。

Megratron是NVIDIA提出的一种分布式训练大规模语言模型的架构，针对Transformer进行了专门的优化，主要采用的是模型并行的方案。这篇文章将描述幻方AI对于NVIDIA Megatron在萤火二号平台上运行的一些实验，以及与我们目前的方法的对比。 ​

模型：GPT

代码：GitHub - NVIDIA/Megatron-LM: Ongoing research training transformer models at scale

环境：幻方萤火二号，16个节点共128张A100（A100-40GB x128）

---

Megatron简介

Megatron是NVIDIA提出的一种由于分布式训练大规模语言模型的架构，针对Transformer进行了专门的优化（也就是大矩阵乘法）。

第一篇论文发表于2019年9月：Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism，主要提出了通过将矩阵分块提高并行度的方法。

第二篇论文发表于2021年4月：Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM，对于分布式中的一些重要的设计，如tensor parallel、pipeline parallel、micro batch size等进行了一些分析与讨论。同时提出了更加精细的pipeline结构与communication模式。

Megatron作者提供的性能结果如下：

 这些测试是基于DGX A100-80GB集群完成的。和萤火二号的测试环境相比，硬件环境上主要存在着如下区别：

• DGX配备了NVLINK、NVSWITCH与IB，使得AllReduce的效率非常高。萤火二号的通讯效率相对更低，尤其是在不使用hfreduce的前提下。
• Megatron使用了A100-80GB而萤火二号使用了A100-40GB。除了显存大小不同之外，二者的内存带宽也不同（2.0T/s vs 1.6T/s）。

除此之外，除去实现细节上的差异，Megatron和我们的方法的主要不同在于

• Megaton支持tensor parallel，并相应地优化了数据传输。
• Megatron对于pipeline机制进行了一些优化（原论文Fig.4）。
• Megatron自行实现了一套DDP框架，而非使用pytorch提供的DDP。
• Megatron自行实现了一些fused kernels，但是为了公平对比被我们disable掉了。

我们关心的问题是，在萤火平台上，Megatron的架构能够达到怎样的训练效率。

我们用于测试的模型配置为：

• hidden size = 3072
• layers = 32
• attention heads = 32
• batch size = 512
• context size = 1024
• vocabulary size = 50264
• micro batch size = 4