在GPU集群上使用Megatron-LM进行高效的大规模语言模型训练

摘要

大型语言模型在多个任务中已取得了最先进的准确率。然而，训练这些模型的效率仍然面临挑战，原因有二：a) GPU内存容量有限，即使在多GPU服务器上也无法容纳大型模型；b) 所需的计算操作数量可能导致不现实的训练时间。因此，提出了新的模型并行方法，如张量并行和流水线并行。不幸的是，简单地使用这些方法在数千个GPU上会导致扩展性问题。本文展示了如何将张量并行、流水线并行和数据并行组合起来，实现数千个GPU的扩展。我们提出了一种新颖的交错流水线调度方法，通过这种方法，可以在与现有方法相当的内存占用下，提升10%以上的吞吐量。我们的方法使得在3072个GPU上，以502 petaFLOP/s的速度对拥有1万亿参数的模型进行训练（每个GPU的吞吐量为理论峰值的52%）成为可能。

1. 引言

基于Transformer的语言模型 [13, 27, 33–35, 42, 46] 在自然语言处理（NLP）领域近年来推动了快速的进展，随着大规模计算的普及和数据集的增大，模型的表现也大幅提升。最近的研究 [11, 40] 显示，大型语言模型在零样本或少样本学习中非常有效，并且在许多NLP任务和数据集上都能达到高准确率。这些大型语言模型具有许多令人兴奋的下游应用，例如客户反馈摘要、自动对话生成、语义搜索和代码自动补全 [1, 4, 5]。因此，最先进的