Megatron-LM：利用模型并行训练十亿级参数语言模型
计算与语言
2020-03-17 v4
摘要
近期语言建模方面的研究表明，训练大型transformer模型可推进自然语言处理应用的技术水平。然而，由于内存限制，非常大型的模型训练起来相当困难。在本工作中，我们提出了训练超大型transformer模型的技术，并实现了一种简单、高效的层内模型并行方法，使训练具有数十亿参数的transformer模型成为可能。我们的方法不需要新的编译器或库改动，与流水线模型并行正交且互补，并且可以通过在原生PyTorch中插入少量通信操作来完全实现。我们通过使用512个GPU收敛高达83亿参数的基于transformer的模型来说明该方法。与维持39 TeraFLOPs(峰值FLOPs的30%)的强单GPU基线相比，我们在整个应用中保持15.1 PetaFLOPs，扩展效率达76%。为证明大型语言模型能进一步推进技术水平(SOTA)，我们训练了一个类似于GPT-2的83亿参数transformer语言模型和一个类似于BERT的39亿参数模型。我们表明，在类BERT模型中仔细处理层归一化的位置对于在模型规模增长时获得更高性能至关重要。使用GPT-2模型，我们在WikiText103(困惑度10.8，对比SOTA的15.8)和LAMBADA(准确率66.5%，对比SOTA的63.2%)数据集上取得了SOTA结果。我们的BERT模型在RACE数据集上取得了SOTA结果(准确率90.9%，对比SOTA的89.4%)。
引用
@article{arxiv.1909.08053,
title = {Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism},
author = {Mohammad Shoeybi and Mostofa Patwary and Raul Puri and Patrick LeGresley and Jared Casper and Bryan Catanzaro},
journal= {arXiv preprint arXiv:1909.08053},
year = {2020}
}