随着近几年 AI 的迅速进步，AI 工程也变得越来越庞大。作为 AI 发展核心之一的大型神经网络也是挑战最大的研究方向之一，大型神经网络的训练难度，则随着集群和模型的不断增长而增强。
这些庞大的训练任务往往需要许多 GPU 的共同协作完成。怎样安排一组 GPU 协同工作，才能更高效地完成复杂的大型神经网络模型的训练任务呢？
随着不同机构的研究人员对大型神经网络越来越深入的研究，已经积累了许多经验；尤其是可以在多个 GPU 上进行并行化模型训练，并应用了越来越多的技巧。
最近，来自 OpenAI 的翁莉莲和格雷格·布罗克曼（Greg Brockman）等研究人员就将这些经验进行了总结。大型神经网络训练时，他们把在多个 GPU 上进行并行化模型训练的安排方式大致分为四类。
6 月 9 日，相关文章以《大型神经网络的训练技巧》（Techniques for Training Large Neural Networks）为题发表在 OpenAI 官网上。
图｜相关文章（来源：OpenAI）
根据该研究内容，尽管在多个 GPU 上进行“并行化”模型训练听上去有些令人生畏、难以理解，但读者只需通过对计算结构在头脑里进行一些构想，就可以更清晰地理解这些技术。事实上，并行化模型训练所做的，就如同在一个人从在繁杂的交通网络中想要从起点 A 到达终点 B，在这一过程中不时地在熙熙攘攘的车流中换乘列车和跑道。
在不使用并行化技术的模型训练中，神经网络的训练过程实际上是一个迭代优化过程。每一个迭代过程中，模型会向前传递一层，也就是计算一批数据中每个训练示例的输出。之后模型则会往回走，“向后”计算一层。
而这一系列过程的各种数据，包括各种参数以及每个参数的优化状态，都会被继续作为输入传递给 Adam 等优化算法中，进行下一次迭代中的参数计算。随着多轮更新和训练数据，该模型也不断地升级，并可得到越来越准确的输出。
图｜以三层模型为例的四种不同并行策略的图示。每种颜色代表一层，虚线分割表示不同的 GPU（来源：OpenAI）按照训练过程中对于GPU的不同分配方式，OpenAI将并行技术分为以下四类，分别是：
数据并行，即不同的 GPU 可运行相同数据的不同子集；
流水线并行，也就是 GPU 各自运行模型的不同层；
张量并行，即将各种数学运算分解到不同的 GPU，举例来说，将一个矩阵乘法拆分至多个 GPU 完成；
专家混合，也就是每个示例都只是通过每一层的某个小部分来处理。
下面对这四种并行方法分别进行简单的解释：
在数据并行中，研究人员将同一参数在多个 GPU 复制，这样在单个 GPU 中分配不同的示例同步处理。在这四种并行策略中，只有数据并行这一策略依旧要求单个 GPU 的内存必须能承载整个模型的大小。尽管这一策略牺牲了内存，需要重复存储同样的副本，但是它带来的好处是可以有效理由多个 GPU 的算力。目前，已经有其他方法可以有效节省内存，例如将参数临时卸载等。
在流水线并行策略中，整个模型被划分到多个 GPU 上，单个 GPU 上只需要存储多有参数的一部分。所以，相比数据并行策略或者不使用并行来说，对于同样大小的模型，流水线并行中每个 GPU 所需的内存大小随着 GPU 数量的增长相应成比减少。
尽管大型模型的拆分并不难，但是由于每一层（不同 GPU）的输入与输出之间存在顺序上的依赖关系。在一个 GPU 的工作开始需要等待前一个 GPU 的输出，如果单纯等待将会浪费许多时间，这些 GPU 在等待的时间也被称为“时间气泡”。
图｜简单的流水线并行策略的示意图，其中整个模型被分为 4 部分并分配在 4 个 GPU 处理。图中 F、B 和 U 分别代表前向、后向和更新。不同 GPU 工作之间的顺序依赖关系导致产生了“时间气泡”（来源：OpenAI）
为了减少“时间气泡“，一些研究人员通过采用数据并行策略中的思路解决该问题。也就是说，让每个 GPU 每次只处理一批数据中的一个子集，让下一个 GPU 的计算时间提前，进而减少整体的等待时间。
图｜减少“时间气泡”的方法图示（来源：OpenAI）
与流水线并行策略将将模型进行逐层地垂直拆分不同，张量并行策略将模型在同一层内进行“水平”拆分。这一策略被认为适用于例如转换器（Transformer）等的许多现代模型，这些模型的训练中最大的计算难点在于，进行一个大型批处理矩阵和大型权重矩阵的相乘。
在张量并行策略中，矩阵的乘法作为成对的行与列的点积运算，把权重矩阵拆分成多个“碎片”，每个碎片交给不同的 GPU 来处理，并使用同样的碎片进行整个矩阵乘积的相关部分的运算，最后再将所有计算结果进行通信组合，进而得出结果。
在专家混合并行策略中，整个网络中只有一部分被用于计算输出。例如，在有多组权重的模型训练中，通过专家混合并行策略可以在计算成本没有增加的前提下使用更多参数，而不同的权重可以交给不同的 GPU 上来处理。
