资讯 人工智能学术
此为临时链接,仅用于文章预览,将在时失效

OpenAI:训练大型神经网络的四种基本方法

作者:我在思考中
2022/06/14 10:43
OpenAI:训练大型神经网络的四种基本方法

来源 | OpenAI

编译 | 黄楠

编辑 | 陈彩娴

大型神经网络是当前人工智能领域的热门话题之一,那么,如何训练大模型?

最近,曾推出大规模预训练模型 GPT-3 的 OpenAI 发表了一篇博文,介绍了基于 GPU 的四种节省内存的并行训练方法,分别是:

OpenAI:训练大型神经网络的四种基本方法

图注:三层模型上各种并行策略,每种颜色代表一层,虚线分隔不同的 GPU。



1

数据并行

「数据并行训练」意味着将相同的参数复制到多个 GPU(通常称为“workers”),并为每个 GPU 分配不同的示例以同时处理。

单单的数据并行要求模型匹配单个 GPU 内存,但当你利用多个 GPU 计算时,代价是存储参数的多个副本。不过,话虽如此,有一些策略可以增加 GPU 可用的有效 RAM,例如,在两次使用之间,可将参数暂时卸载到 CPU 内存。

随着每次数据并行 worker 更新其参数副本,它们需要相互协调,以确保每个 worker 都继续具有相似的参数。最简单的方法是在 worker 之间引入「阻塞通信」:

步骤 1:独立计算每个worker上的梯度;

步骤 2:将不同 worker 的梯度平均;

步骤 3:在每个 worker 上独立计算相同的新参数。

步骤 2 是一个阻塞平均值,它需要传输大量数据(与 worker 数量乘以参数大小成正比),这可能会损害训练的吞吐量。有各种异步同步方案可以消除这种损耗,但会损害学习效率;因此在实践中,人们普遍坚持同步方法。



2

流水线并行

在流水线并行训练中,研究者会将模型的顺序块划分到 GPU 上,每个 GPU 只保存一小部分参数,因此,相同模型的每个 GPU 消耗的内存按比例减少。

将大型模型拆分为连续层的块很简单,但由于层的输入和输出之间存在顺序依赖关系,因此,在 worker 等待前一台机器的输出用作其输入时,一个幼稚的执行可能会导致出现大量空闲时间。这些等待时间块被称为「泡沫」(bubbles),即浪费了本可以由空闲机器来完成的计算。

OpenAI:训练大型神经网络的四种基本方法

图注:一个简单的流水线并行设置插图,其中,模型被垂直分成 4 个分区。worker 1 主持第一层的模型参数(最接近输入),而 worker 4 主持第 4 层(最接近输出)。“F”、“B”和“U”分别代表前向、后向和更新操作。下标会指示在哪个 worker 上运行操作。由于顺序依赖性,数据一次由一个 worker 处理,导致产生了大量的空闲时间“泡沫”。

我们可以重用数据并行的想法,通过让每个 worker 一次只处理数据元素的一个子集,来降低产生时间泡沫的成本,从而使我们能巧妙地将新计算与等待时间重叠。核心思想是,将一个批次拆分为多个微批次,每个微批次的处理速度都应该成比例地加快,并且每个 worker 在下一个微批次可用时立即开始工作,从而加快管道执行。有了足够的微批次, worker 可以在大部分时间被利用,并且在步骤开始和结束时「泡沫」最小。梯度在微批次之间进行平均,并且只有在所有微批次完成后才会更新参数。

模型拆分的 worker 数量通常称为「管道深度」(pipeline depth)。

在前向传递期间,worker 只需将其层块的输出(称为「激活」)发送给下一个 worker;在反向传递期间,它仅将这些激活的梯度发送给前一个工作人员。如何安排这些通道以及如何跨微批次聚合梯度有很大的设计空间。例如,方法 GPipe 是让每个工作进程连续向前和向后传递,然后在最后同步聚合来自多个微批次的梯度;而 PipeDream 会安排每个 worker 交替处理的前向和后向通道。

OpenAI:训练大型神经网络的四种基本方法

OpenAI:训练大型神经网络的四种基本方法

图注:GPipe 和 PipeDream 流水线方案的比较,每批使用 4 个微批次。微批次 1-8 对应于两个连续的数据批次。图中“number”表示在哪个微批次上操作,下标标记 worker  ID。注意,PipeDream 通过使用陈旧参数执行一些计算来获得更高的效率。



3

张量并行

管道并行性将模型逐层“垂直”拆分,也可以在一个层内“水平”拆分某些操作,这通常称为张量训练。

对于许多现代模型(例如Transformer),计算瓶颈是将激活批处理矩阵与大权重矩阵相乘。矩阵乘法可以认为是成对的行和列之间的点积;可以在不同的 GPU 上计算独立的点积,或者在不同的 GPU 上计算每个点积的部分并总结结果。无论采用哪种策略,我们都可以将权重矩阵分割成大小均匀的“碎片”,将每个碎片托管在不同的 GPU 上,并使用该碎片计算整个矩阵乘积的相关部分,然后再进行通信以组合结果。

一个例子是Megatron-LM,它在 Transformer 的自注意力和 MLP 层内并行化矩阵乘法。PTD-P使用张量、数据和流水线并行,其流水线调度为每个设备分配了多个不连续的层,以增加网络通信为代价来减少泡沫损耗。

有时,网络输入可以跨维度并行化,相对于交叉通信具有高度的并行计算。序列并行就是这样一种想法,其中输入序列在时间上被分成多个子示例,通过允许计算继续进行更细粒度的示例,来按比例减少峰值内存消耗。



4

专家混合 (MoE)

使用专家混合(MoE)方法,只有小部分网络用于计算任何一个输入的输出。

一个示例方法是拥有多组权重,并且网络可在推理时通过门控机制选择要使用的权重组,这能在不增加计算成本的情况下启用更多参数。每组权重都被称为“专家”,且希望网络能学会为每个专家分配专门的计算和技能。不同的专家可以主持不同的 GPU ,从而提供了一种明确的方式来扩大用于模型的 GPU 数量。

OpenAI:训练大型神经网络的四种基本方法

图注:门控网络只选择了n个专家中的2个。

GShard 将 MoE Transformer 的参数扩展到 6000 亿个参数,其中仅将 MoE 层拆分到多个 TPU 设备上,其他层则完全复制。Switch Transformer 通过将一个输入路由到单个专家,将模型大小扩展到数万亿个参数,具有更高的稀疏性。



5

其他节省内存的设计

还有许多其他的计算策略,可以使训练越来越大的神经网络更容易处理。例如:

要计算梯度,需要保存原始激活,这会消耗大量设备 RAM。检查点(也称为激活重新计算)存储激活的任何子集,并在反向传递期间,及时重新计算中间的激活,以最多一个额外完整前向传递的计算成本,节省了大量内存。人们还可以通过选择性激活重新计算,来不断权衡计算和内存成本,这是对激活的子集进行检查,其存储成本相对较高,但计算成本较低。

混合精度训练是使用较低精度的数字(最常见的是FP16)来训练模型。现代加速器可以使用较低精度的数字达到更高的 FLOP 计数,并且还能节省设备 RAM。在适当的照顾下,产生的模型几乎可以不损失任何精度。

卸载是将未使用的数据临时卸载到 CPU 或不同设备之间,在需要时将其读回。幼稚的执行会大大减慢训练速度,但复杂的实现方式会预先获取数据,使设备永远不需要等待。这个想法的一个实现是ZeRO,它可将参数、梯度和优化器状态分割到所有可用的硬件上,并根据需要将它们具体化。

Memory Efficient Optimizers已经提出了内存效率优化器,以减少优化器所维护的运行状态的内存占用,例如Adafactor。

压缩也可用于存储网络中的中间结果。例如,Gist压缩为后向传递而保存的激活;DALL-E在同步梯度之前压缩梯度。

参考链接:

https://openai.com/blog/techniques-for-training-large-neural-networks/
OpenAI:训练大型神经网络的四种基本方法

雷峰网(公众号:雷峰网)

长按图片保存图片,分享给好友或朋友圈

OpenAI:训练大型神经网络的四种基本方法

扫码查看文章

正在生成分享图...

取消
相关文章