梯度|北大校友“炼丹”分享：OpenAI如何训练千亿级模型？

文章插图

编译 | 琰琰
大规模深度神经网络训练仍是一项艰巨的挑战，因为动辄百亿、千亿参数量的语言模型，需要更多的 GPU 内存和时间周期。这篇文章从如何多GPU训练大模型的角度，回顾了现有的并行训练范式，以及主流的模型架构和内存优化设计方法。
本文作者Lilian Weng现为OpenAI应用人工智能研究负责人，主要从事机器学习、深度学习和网络科学研究。她本科毕业于香港大学，硕士就读于北京大学信息系统与计算机科学系，之后前往印度安纳大学布鲁顿分校攻读博士。
Lilian Weng经常在个人博客分享学习和工作笔记，感兴趣的可以戳这里：
https://lilianweng.github.io/lil-log/。
“炼大模型”已成为人工智能领域的主流研发趋势。从GPT-3的1750亿，到如今悟道2.0的1.75万亿，超大语言模型在 NLP 基准任务中不断刷新SOTA。
然而，参数和数据集的快速增长让 GPU 算力开始捉襟见肘。单个GPU内存已经远远不能满足大模型的需求。如，阿里用480块GPU训练千亿模型；英伟达用3072块GPU训练万亿模型；谷歌用2048块TPU训练1.6万亿模型（1 TPU约等于2～3 GPU）。
如何利用上百块GPU上训练大规模语言模型？并行计算是一种行之有效的方法。
近日，OpenAI 研究员Lilian Weng分享干货文章，从并行训练（数据并行、模型并行、管道并行、张量并行）、混合专家、内存节省设计（CPU卸载、重新激活计算、混合精度训练、高效存储优化器）三个方面回顾了现阶段多GPU训练大模型的主流方法。

文章插图

AI科技评论这篇文章编译如下，供想成为“炼丹师”的朋友参考。

并行训练

大规模神经网络模型对存储空间有强烈的需求，单个GPU的内存已经远远不够。究其原因，一方面模型权重有数百亿个浮点数，随机梯度下降和Adam优化需要极高的计算成本；另一方面在预训练阶段，大模型与大规模语料库的配对需要很长的时间周期。综合来看，跨GPU并行计算显得尤为重要。
并行计算在数据、模型架构和张量等不同维度上都可以操作，接下来本文将具体介绍一些主流方法：
数据并行
数据并行（ Data parallelism ，DP）最简单的方法是将相同的模型权重复制到worker节点，并分配一部分数据以同时进行处理。我们知道，如果模型的参数量大于单个GPU节点的内存，DP无法正常工作，GeePS架构（Cui等人，2016）的解决思路是使用有限的GPU内存。也就是，如果模型太大无法嵌入到一台机器，就将暂时未使用的参数卸载回CPU。
数据交换传输通常在后端进行（不干扰训练计算），在每个Mini-batch计算结束后，worker需要同步梯度或权重，以保证学习效率。现有的同步方法有两种，各自优缺点如下：
1、批量同步并行（BSP）：worker在每个Mini-batch结束时同步数据，这种方法保证了模型权重传递的及时性，但每台机器都必须排队等待其他机器发送梯度。
2、异步并行（ASP）：每个GPU采用异步方式处理数据，这种方法避免了不同机器之间的相互等待或暂停，但影响了权重传递的时效，降低了统计学习效率。而且即使增加计算时长，也不会加快训练的收敛速度。
在中间某些地方的每一次迭代（＞1）都需要同步全局梯度。自Pytorch v1.5版（Li等人，2021年）提出后，该特征在分布式数据并行（Distribution Data Parallel，DDP）中被称为“梯度累积（gradient accumulation）”。分桶梯度(bucketing gradients)避免立即执行AllReduce操作，而是将多个梯度存储到一个AllReduce中以提高吞吐量，并基于计算图优化计算和通信调度。