OpenAI 研究员最新博客：如何在多GPU上训练真正的大模型？

【专栏：前沿进展】近年来，在大规模预训练语言模型的帮助下，许多NLP模型在基准测试任务中取得了更好的结果。如何训练大而深的神经网络是一个挑战，需要大量的GPU内存和很长的训练时间。本文回顾了几种流行的并行训练范例，以及各种模型结构和内存节省设计，使跨大量GPU训练大型神经网络成为可能。

然而，单个GPU卡的内存有限，许多大模型的大小已经超过了单个GPU，目前，为解决此类问题，训练深且大的神经网络的主要方法有训练并行加速、各种模型架构以及内存节省设计等。常见的并行加速方法有以下几种：数据并行性、模型并行性、流水线并行以及张量平行。模型架构方面主要有专家混合（MoE）方法。此外，还有其他节省内存的设计方法，如：CPU卸载、激活重新计算、混合精度训练、压缩以及内存高效优化器等等。

本文编译自OpenAI科学家Lilian Weng（翁荔）最新博客文章（lilianweng.github.io/lil-log/2021/09/24/train-large-neural-networks.html)，智源社区已经获得翻译授权

作者：翁荔

编译：马瑞军

校对：李梦佳

01

训练并行性

训练大模型的主要瓶颈是对大量GPU内存的强烈需求，远远高于单个GPU上的内存。大模型在训练过程中，GPU内存主要用在模型权重参数存储（例如数百亿个浮点数），中间计算输出存储（例如梯度和优化器状态，例如Adam中的动量和变化），这使得GPU成本十分昂贵。除此之外，训练一个大模型通常时通常需要与一个大型训练语料库配对，因此单个过程可能需要很长时间。

因此，并行性是必要的，同时，并行可以使用在不同的维度上，包括数据、模型架构和张量操作等。

数据并行性

数据并行性（Data parallelism (DP)）最简单的方法是将相同的模型权重复制到多个worker中，并将一部分数据分配给每个worker以同时进行处理。

如果模型规模大于单个GPU的内存，Naive DP无法正常工作时。GeePS（Cui 等人，2016 年）之类的方法将暂时未使用的参数卸载回 CPU，以使用有限的 GPU 内存。数据交换传输在后端进行，且不干扰训练计算。

在每个小批量结束时，workers需要同步梯度或权重，以替换旧参数。常见有两种主要的同步方法，它们都有明确的优缺点：

1）大容量同步并行（ Bulk synchronous parallels (BSP)）：workers在每个小批量结束时同步数据。这种方法可以防止模型权重过时，同时获得良好的学习效率，但每台机器都必须停止并等待其他机器发送梯度。

2）异步并行（Asynchronous parallel (ASP)）：每个GPU工作进程异步处理数据，无需等待或暂停。然而，这种方法很容易导致网络使用陈旧的权重参数，从而降低统计学习效率。即使它增加了计算时间，也可能不会加快收敛的训练时间。

中间的某个地方是在每次x迭代时，全局同步梯度（x＞1）。自Pytorch v1.5版（Li等人，2021年）以来，该特征在平行分布数据（DDP）中被称为“梯度累积”。Bucket 梯度计算方法避免了梯度的立即AllReduce，而是将多个梯度变化值存储到一个AllReduce中以提高吞吐量，可以基于计算图进行计算和通信调度优化。

 

 图1：Pytorch DDP的伪代码。（图片来源：李等人，2021年)

模型并行性

模型并行性（Model parallelism: MP）目的是解决模型权重不能适应单个节点的情况，通过将计算和模型参数分布在多台机器上进行训练。在数据并行中，每个worker承载整个模型的完整副本，而在MP中，每个worker上只分配模型参数的一小部分，从而减少了内存使用和计算。

由于深度神经网络通常包含一堆垂直层，因此将一个大型模型逐层拆分感觉很简单，其中一组连续的小层被分组到一个工作层上的一个分区中。然而，通过多个具有顺序依赖性的工作线程来运行每个数据批，会导致大量的等待时间和计算资源利用率低下的问题。

 

图2：一个简单的模型并行设置，其中模型垂直分割为4个分区。由于顺序依赖性，数据一次由一个工作者处理，从而导致训练过程中出现大量“间隔时间”。（图片来源：黄等人，2019年)

流水线并行

通道并行（Pipeline parallelism: PP）将模型并行与数据并行相结合，以减少部分训练过程中出现的空闲时间。其主要思想是将一个小批量拆分为多个微批次，并使worker在每个阶段中能够同时处理一个微批次。需要注意的是，每个微批次需要两次传递，一次向前，一次向后。worker之间的通信仅传输激活（向前）和梯度（向后）。这些通道的调度方式以及梯度的聚合方式在不同的方法中有所不同。分区（workers）的数量也称为通道深度。

在GPipe（Huang et al.2019）中，来自多个微批次的梯度在最后聚合并同步应用。同步梯度下降的操作，与workers数量无关，可以保证学习的一致性和效率。如图3所示，“间隔时间”的情况仍然存在，但比图2中的少得多。设定m为平均分割微批次和d为分区，假设每个微批次的向前和向后都需要一个单位的时间，“间隔时间”的分数为：

GPipe论文观察到，如果微组的数量超过分区数量的4倍（m>4d）（当应用激活重新计算时），则“间隔时间