作者|Lilian Weng、Greg Brockman
翻译|董文文
AI 畛域的许多最新进展都围绕大规模神经网络开展，但训练大规模神经网络是一项艰巨的工程和钻研挑战，须要协调 GPU 集群来执行单个同步计算。
随着集群数和模型规模的增长，机器学习从业者开发了多项技术，在多个 GPU 上进行并行模型训练。
乍一看，这些并行技术令人生畏，但只需对计算构造进行一些假如，这些技术就会变得清晰——在这一点上，就像数据包在网络交换机之间传递一样，那也只是从 A 到 B 传递并不通明的位（bits）。
三层模型中的并行策略。每种色彩代表一层，虚线分隔不同的 GPU。
训练神经网络是一个迭代的过程。在一次迭代过程中，训练数据通过模型的 layer（层）进行前向传递，对一批数据中的训练样本进行计算失去输入。而后再通过 layer 进行反向传递，其中，通过计算参数的梯度，能够失去各个参数对最终输入的影响水平。
批量均匀梯度、参数和每个参数的优化状态会传递给优化算法，如 Adam，优化算法会计算下一次迭代的参数 (性能更佳）并更新每个参数的优化状态。随着对数据进行屡次迭代训练，训练模型会一直优化，失去更加准确的输入。
不同的并行技术将训练过程划分为不同的维度，包含：
- 数据并行（Data Parallelism）——在不同的 GPU 上运行同一批数据的不同子集；
- 流水并行（Pipeline Parallelism）——在不同的 GPU 上运行模型的不同层；
- 模型并行（Tensor Parallelism）——将单个数学运算（如矩阵乘法）拆分到不同的 GPU 上运行；
- 专家混合（Mixture-of-Experts）——只用模型每一层中的一小部分来解决数据。
本文以 GPU 训练神经网络为例，并行技术同样也实用于应用其余神经网络加速器进行训练。作者为 OpenAI 华侨工程师 Lilian Weng 和联结创始人 & 总裁 Greg Brockman。
1
数据并行
数据并行是指将雷同的参数复制到多个 GPU 上，通常称为“工作节点（workers）”，并为每个 GPU 调配不同的数据子集同时进行解决。
数据并行须要把模型参数加载到单 GPU 显存里，而让多个 GPU 计算的代价就是须要存储参数的多个正本。话虽如此，还有一些办法能够减少 GPU 的 RAM，例如在应用的间隙长期将参数卸载（offload）到 CPU 的内存上。
更新数据并行的节点对应的参数正本时，须要协调节点以确保每个节点具备雷同的参数。
最简略的办法是在节点之间引入阻塞通信：（1）独自计算每个节点上的梯度；(2) 计算节点之间的均匀梯度；(3) 独自计算每个节点雷同的新参数。其中，步骤 (2) 是一个阻塞平均值，须要传输大量数据（与节点数乘以参数大小成正比），可能会侵害训练吞吐量。
有一些异步更新计划能够打消这种开销，然而会侵害学习效率；在实践中，通常会应用同步更新办法。
2
流水并行
流水并行是指按程序将模型切分为不同的局部至不同的 GPU 上运行。每个 GPU 上只有局部参数，因而每个局部的模型耗费 GPU 的显存成比例缩小。
将大型模型分为若干份间断的 layer 很简略。然而，layer 的输出和输入之间存在程序依赖关系，因而在一个 GPU 期待其前一个 GPU 的输入作为其输出时，奢侈的实现会导致呈现大量闲暇工夫。这些闲暇工夫被称作“气泡”，而在这些期待的过程中，闲暇的机器本能够持续进行计算。
一个奢侈的流水并行设置，其中模型按 layer 垂直分成 4 个局部。worker 1 托管网络第一层（离输出最近）的模型参数，而 worker 4 托管第 4 层（离输入最近）的模型参数。“F”、“B”和“U”别离代表前向、反向和更新操作。下标批示数据在哪个节点上运行。因为程序依赖性，数据一次只能在一个节点上运行，从而会导致大量闲暇工夫，即“气泡”。
为了缩小气泡的开销，在这里能够复用数据并行的打法，核心思想是将大批次数据分为若干个微批次数据（microbatches），每个节点每次只解决一个微批次数据，这样在原先期待的工夫里能够进行新的计算。
每个微批次数据的处理速度会成比例地放慢，每个节点在下一个小批次数据开释后就能够开始工作，从而放慢流水执行。有了足够的微批次，节点大部分工夫都在工作，而气泡在过程的结尾和完结的时候起码。梯度是微批次数据梯度的平均值，并且只有在所有小批次实现后才会更新参数。
模型拆分的节点数通常被称为流水线深度（pipeline depth）。
在前向传递过程中，节点只需将其 layer 块的输入（激活）发送给下一个节点；在反向传递过程中，节点将这些激活的梯度发送给前一个节点。如何安顿这些过程以及如何聚合微批次的梯度有很大的设计空间。GPipe 让每个节点间断前向和后向传递，在最初同步聚合多个微批次的梯度。PipeDream 则是让每个节点交替进行前向和后向传递。
GPipe 和 PipeDream 流水计划比照。每批数据分为 4 个微批次，微批次 1 - 8 对应于两个间断大批次数据。图中，“（编号）”示意在哪个微批次上执行操作，下标示意节点 ID。其中，PipeDream 应用雷同的参数执行计算，能够取得更高的效率。
3
模型并行
在流水并行中，模型沿 layer 被“垂直”拆分，如果在一个 layer 内“程度”拆分单个操作，这就是模型并行。许多古代模型（如 Transformer）的计算瓶颈是将激活值与权重相乘。
矩阵乘法能够看作是若干对行和列的点积：能够在不同的 GPU 上计算独立的点积，也能够在不同的 GPU 上计算每个点积的一部分，而后相加失去后果。
无论采纳哪种策略，都能够将权重矩阵切分为大小平均的“shards”，不同的 GPU 负责不同的局部。要失去残缺矩阵的后果，须要进行通信将不同局部的后果进行整合。
Megatron-LM 在 Transformer 的 self-attention 和 MLP layer 进行并行矩阵乘法；PTD- P 同时应用模型、数据和流水并行，其中流水并行将多个不间断的 layer 调配到单设施上运行，以更多网络通信为代价来缩小气泡开销。
在某些场景下，网络的输出能够跨维度并行，绝对于穿插通信，这种形式的并行计算程度较高。如序列并行，输出序列在工夫上被划分为多个子集，通过在更细粒度的子集上进行计算，峰值内存耗费能够成比例地缩小。
4
混合专家（MoE）
混合专家（MoE）模型是指，对于任意输出只用一小部分网络用于计算其输入。在领有多组权重的状况下，网络能够在推理时通过门控机制抉择要应用的一组权重，这能够在不减少计算成本的状况下取得更多参数。
每组权重都被称为“专家（experts）”，现实状况是，网络可能学会为每个专家调配专门的计算工作。不同的专家能够托管在不同的 GPU 上，这也为扩充模型应用的 GPU 数量提供了一种明确的办法。
混合专家（MoE）层。门控网络只抉择了 n 个专家中的 2 个（图片改编自：Shazeer et al., 2017）。
GShard 将 MoE Transformer 扩大到 6000 亿个参数，其中 MoE layers 被拆分到多个 TPU 上，其余 layers 是齐全反复的。Switch Transformer 将输出只路由给一个专家，将模型大小扩大到数万亿个参数，具备更高的稠密性。
5
其余节俭内存的设计
除了以上的并行策略，还有很多其余的计算策略能够用于训练大规模神经网络：
- 要计算梯度，须要保留原始激活值，而这会耗费大量设施显存。Checkpointing（也称为激活重计算）存储激活的任何子集，并在反向流传时及时从新计算两头的激活。这能够节俭大量内存，而计算成本最多就是减少一个残缺的前向传递。还能够通过 选择性激活重计算（https://arxiv.org/abs/2205.05198）在计算和内存老本之间一直衡量，也就是对那些存储老本绝对较高但计算成本较低的激活子集进行查看。
- 混合精度训练（https://arxiv.org/abs/1710.03740）是应用较低精度的数值（通常为 FP16）来训练模型。古代加速器能够用低精度的数值实现更高的 FLOP 计数，同时还能够节俭设施显存。解决切当的话，简直不会损失生成模型的精度。
- Offloading是将未应用的数据长期卸载到 CPU 或其余设施上，在须要时再将其读回。奢侈实现会大幅升高训练速度，而简单的实现会预取数据，这样设施不须要再期待数据。其中一个实现是 ZeRO（https://arxiv.org/abs/1910.02054），它将参数、梯度和优化器状态宰割到所有可用硬件上，并依据须要将它们实现。
- 内存效率优化器 可缩小优化器保护的运行状态的内存，例如 Adafactor。
- 压缩 可用于存储网络的两头后果。例如，Gist 能够压缩为反向传递而保留的激活；DALL·E 能够在同步梯度之前压缩梯度。
（原文：https://openai.com/blog/techn…）
欢送下载体验 OneFlow v0.7.0 最新版本：
https://github.com/Oneflow-In…