Megatron-LM论文阅读
最近对于大模型的研究越来越多，笔者也对一些大模型的技术报告进行阅读，但是感觉对于大模型的认识仍然不够清晰，究其原因是没有自己亲手走一下训练大模型的流程，所以想要先对训练大模型的框架进行学习，然后等到有机会之后（有显卡）再上手训练一个大模型。经过调研之后，笔者打算首先从Megatron-LM框架开始学习，于是选择了这篇Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism论文进行阅读。
论文地址和相关源码链接如下：
论文地址：https://arxiv.org/abs/1909.08053
相关源码：https://github.com/NVIDIA/Megatron-LM
摘要
由于内存限制，大型Transformer模型可能难以训练，论文提出了训练超大型Transformer模型的技术，并实现了一种简单、有效的层内模型并行方法，使训练具有数十个亿参数的Transformer模型成为可能，论文提出的方法只需在原生pytorch中插入少量通信操作即可完全实现。论文使用512个GPU将基于变压器的模型收敛到 83 亿个参数，以此说明这种方法。 与可维持39TeraFLOPs（峰值FLOPs的30%）的强大单GPU基线相比，我们在整个应用中维持了15.1PetaFLOPs，扩展效率高达76%。
\[ 512 \times 39 = 19,968 \ TFLOPS = 19.968 \ PFLOPS \] \[ 15.1 / 19.968 \approx 0.76 \]
另外，论文也证明了大型语言模型可以进一步提升SOTA。论文训练了一个类似于GPT-2的83亿参数变压器语言模型和一个类似于BERT的39亿参数模型，利用GPT-2模型在WikiText103和LAMBADA数据集上取得了SOTA结果，利用BERT模型在RACE数据集上取得了SOTA结果。
介绍
- 可用计算能力的不断增加和数据的不断丰富使得人们能够通过无监督预训练来训练越来越大的语言模型，经验证据表明，更大的语言模型对于文章补全、问题解答和自然语言推理等NLP任务的作用更大，通过在下游自然语言任务中对预训练的语言模型进行微调，可以获得最先进的结果。
- 随着模型变得越来越大，超出了现代处理器的内存限制。现有的模型并行化方法通过对模型进行分区，使权重及相关优化器状态无需同时驻留在处理器上，从而克服这一限制，例如，
GPipe和Mesh-Tensorflow提供了不同类型的模型并行框架。 不过，它们需要重写模型，并依赖于仍在开发中的定制编译器和框架。 - 论文利用层内模型并行性实现了一种简单高效的模型并行方法。利用基于Transformer的语言模型的固有结构，在PyTorch中高效地实现了简单的模型并行，无需定制C++代码或编译器。
- 论文还对提出方法的可扩展性，以及模型大小对准确性的影响进行分析。
论文的贡献可以总结如下：
- 通过对现有的Pytorch Transformer实现一些有针对性的修改，实现了一种简单而有效的模型并行方法，然后对模型和数据并行技术进行了深入的实证分析，并使用512个gpu展示了高达76%的扩展效率。
- 研究表明，在类似
BERT的模型中，层归一化的位置对于随着模型的增长而提高准确度至关重要。另外，扩大模型规模可提高GPT-2（参数多达83亿）和BERT（参数多达39亿）模型的精确度，它们在测试集上达到了最先进的结果。
背景
在这里只对几个特别需要注意的点进行强调：
GPT-2和BERT的模型架构：GPT-2使用decoder-only架构，而BERT使用encoder-only架构。另外，GPT-2和BERT都使用GeLU非线性和层归一化作为多级注意和前馈层的输入，而原始Transformer使用ReLU非线性并对输出进行层归一化。数据并行和模型并行：
将深度神经网络训练扩展到众多硬件加速器有两个核心范式：数据并行，即在多个worker上分割训练小批量数据；模型并行，即在多个worker上分配模型的内存使用和计算。 通过按可用worker数量成比例增加minibatch的规模（即弱扩展weak scaling），可以观察到训练数据吞吐量接近线性扩展。然而，大批量训练会给优化过程带来复杂性，导致准确率降低或收敛时间延长，从而抵消了训练吞吐量增加带来的好处.