1.1.3 并行模式

分布式机器学习的并行模式描述了计算节点之间应以何种方式进行协作，即计算节点各自分担什么子任务，以及应当交互什么内容。根据用户不同的需求，常用的并行模式包括模型并行、流水线并行、数据并行和混合并行四种。模型并行适用于大模型训练任务；流水线并行则优化了模型并行的算力利用率，进一步提高了并行训练效率；数据并行适用于大数据训练任务。但是在现实应用中，大模型和大数据往往共存，上述三种并行模式通常需要结合使用，它们的结合被称为混合并行。在本节中，我们将分别介绍这四种并行模式。

1.1.3.1 模型并行

模型并行常用于模型参数量巨大的情况。例如，T5-XXL模型拥有11B个可训练参数，需要占用44GB的显存空间，而RTX 3080TI GPU仅有12GB的显存容量，单卡无法容纳规模较大的模型。因此，模型并行将大模型进行切片，将各个较小的模型切片分配给多张计算卡共同存储，并交换切片的输出以打通多卡计算过程。在模型并行的框架下，计算节点之间有较为复杂的依赖关系，某个计算节点的输出是另一个计算节点的输入，于是两个节点之间可建立一条有向边表示中间结果的传输数据流，从而构成分布式计算图。

总的来说，给定一个深度网络模型，可以对其按层切片、按神经元切片以及更为精细的混合切片，如图1-5所示。首先，由于深度网络模型常由多个神经网络层构成，所以最直观的方式就是按层切片，如图1-5（b）所示，前两层网络的切片分配给GPU0，后两层网络的切片分配给GPU1。这种切片方式虽然直观，但GPU1的计算依赖于GPU0的输出，两卡不能并行，算力利用率低，仅能支撑大模型训练且不高效。

第二种方式是按神经元切片，即将各个网络层的神经元及其参数分割给不同的计算卡，如图1-5（c）所示。在此种切片方式下，GPU0和GPU1各拥有一半的相同切片，数据样本同时输入两个切片，并在切片中的各层频繁交换中间结果。与按层切片相比，按神经元切片的计算卡之间可以完全并行。然而，中间结果（如激活图）的尺寸可能很大，频繁的中间结果交换会引入高传输延迟，也容易导致低算力利用率。

第三种混合切片方式结合前两种切片方式的优势，根据深度网络模型各层的参数量、输出激活图的尺寸以及计算卡之间的数据吞吐率，结合使用按层切片和按神经元切片，可对深度网络进行更精细的切片。如图1-5（d）所示，第一、四层网络的切片分配给GPU0，第二、三层网络按神经元进行切片，分别分配给GPU1和GPU2。于是，数据样本首先输入GPU0，其输出作为GPU1和GPU2的输入，GPU1和GPU2在计算过程中交换中间结果，它们的输出又传回GPU0进行最后的计算。通常来说，模型切片越精细，跨GPU通信越复杂，通信代价也就越高。通过谨慎设定切片策略，在通信代价和并行度之间折中，可以有效提升整体训练效率[40-42]。

下面我们以一个真实深度网络AlexNet为例，对模型并行的实现细节做详细阐述。AlexNet包含五个卷积层和三个全连接层，总计参数量61×106个，需占用244MB显存空间。在2012年，Alex等人使用的GTX 580 GPU仅有3GB显存容量，在排除训练数据、中间结果（如模型梯度）以及外接显示器等占用的显存后，可用显存空间相当紧张。于是他们设计了图1-6所示的基于双卡模型并行的AlexNet网络架构，图1-6中的K代表103，表示“个”参数；M代表106，即1M=106。该架构采用类似于图1-5（c）的按神经元切片方式，在双卡的分担下，GPU0仅需存储和训练3.2×107个参数，GPU1为2.8×107个，单卡存储压力得以缓解。不同的是，仅第二、五、六、七层的输出激活图会在计算卡之间交换。当批数据量为128个训练样本时，GPU之间每次迭代只需要交换28.5MB数据，相比原本需要交换的126MB数据减少了约77%，大幅降低了计算卡之间的传输开销。

模型并行解决了大模型训练的问题，但其扩展能力是有限的，一方面取决于模型结构是否可切片以及最大切片数量；另一方面，一味地增大模型并行度并不一定能取得训练加速[42]。在实践中，若数据量不大且多卡传输具有高带宽，可首先考虑用模型并行来训练大模型。

1.1.3.2 流水线并行

虽然模型并行能够训练大模型，但GPU算力资源利用严重不足。其一，按层切片的方式需要串行处理各个切片，同一时间只有一个GPU在处理模型切片，其余GPU则被闲置。其二，若按神经元切片，则各个切片在处理过程中需要频繁交换中间结果，长时间的GPU通信也会使GPU算力闲置。为提高GPU算力利用率，流水线并行结合了模型并行和数据并行。首先，对于模型并行，流水线并行沿用按层切片的模型并行方法，将大模型的不同网络层分配给多个GPU共同承载。其次，数据并行将原本的批量训练数据再次细分为多个小批次，当一个GPU处理完当前小批次并递交下一个GPU后，就立即处理下一个小批次，如此形成流水线式的持续忙碌的工作模式。图1-7（a）和图1-7（b）形象地显示出模型并行和传统任务（或模型推断阶段）的流水线并行工作模式。

首先回顾一个小批次任务的训练过程，考虑图1-8所示简单的模型并行切片示例，全连接神经网络的四个层有相同的参数量，且四个层分别分配给四个GPU。在前向传播过程中，数据样本输入GPU0，经过GPU0至GPU3依次处理后，GPU3计算损失值和末层网络参数的梯度，并依次反向传播给GPU2至GPU0以计算前层网络参数的梯度。由此可知，深度学习模型训练阶段的流水线并行并非图1-7（b）所示简单的单向流水线，而是后序批次的前向传播和前序批次的反向传播同时存在，使之成为图1-7（c）所示的更为复杂的双向流水线，需要特别设计多批次训练任务的调度方案。

本节介绍三种典型的流水线并行方案，分别是PipeDream[43-44]、GPipe[45]和DAPPLE[46]。假设有四个GPU，每个GPU负责一个模型切片，每个模型切片的计算时间是一个单位，以下举例说明三种流水线并行方案。请注意，当某个GPU处理完一个小批次后，传输计算结果的通信过程与处理下一个小批次的计算过程相互无依赖，两个过程可以重叠执行。为了更清晰地展示流水线并行的突出优势，本例中忽略了GPU通信延迟。原始模型并行的调度时间线如图1-9（a）所示，在同一时刻，仅有一个GPU处于工作状态，其余GPU被闲置，GPU算力的平均利用率仅有25%。

第一种方案是PipeDream流水线并行，如图1-9（b）所示。在启动阶段，PipeDream连续启动四个小批次的前向计算，一旦某个GPU开始执行反向计算，后续就切换到前向和反向交替计算的模式。当某个GPU需要处理多个小批次时，如GPU3在第5个单位时间面临第1个小批次的反向计算（B₁）和第2个小批次的前向计算（F₂），PipeDream优先调度最早的小批次（即批次编号最小）。在第10个单位时间时，PipeDream进入稳定阶段，所有GPU持续忙碌，达到100%的算力利用率。但需注意的是，由于小批次完成前向传播和反向传播需要一定的时间，其间模型参数会被多次更新。例如，GPU0在第9个单位时间时启动批次F₅，但直至第16个单位时间才开始计算梯度B₅，而期间模型参数已经被批次2～4更新了三次。这意味着B₅将用更新版本的模型参数计算梯度，这破坏了梯度的计算法则，也会阻碍训练收敛。因此，PipeDream需要存储近期批次的模型参数，确保梯度计算时模型参数的一致性。另外，如何在异构且变化的GPU网络中实现负载均衡的动态模型切片与任务放置[47-48]，以及如何缓解异步流水线中过期参数的影响[49-50]，也是值得关注的问题。

第二种方案是谷歌提出的GPipe流水线并行，如图1-9（c）所示。在前向传播阶段，GPipe连续启动五个小批次的前向计算与传播。在所有小批次的前向传播完成后，进入反向传播阶段，连续启动所有小批次的反向计算与传播。当所有小批次的反向传播完成后，累计求和五个小批次的梯度，并同步更新各个GPU的模型参数。不同于PipeDream的异步更新，GPipe有明确的同步屏障，属于同步流水线模式，模型参数在抵达同步屏障之前不会被更新，保证了同一设备上的前向计算和反向计算都是基于相同的模型参数，所以GPipe可以正确计算并简单累计求和模型梯度。上述前向传播阶段、反向传播阶段和同步更新阶段合称为一个迭代周期。理论上，GPipe的算力利用率可维持在62.5%。虽然GPipe不像PipeDream那样需要存储多个小批次的模型参数，但它在完成反向计算之前不能丢弃前向传播阶段的中间结果，仍需要消耗大量显存来缓存这些中间结果。尽管一些中间结果可以在反向计算时复现，但重新计算这些中间结果又会引入额外20%的计算开销，需要谨慎权衡。

因为反向计算依赖前向计算的输出激活图，GPU需要缓存这些中间结果直至反向计算完成，所以DAPPLE通过提前启动反向计算来减少显存消耗。如图1-9（d）所示，不同于GPipe一次性启动全部小批次任务，DAPPLE先启动少数小批次任务，并严格调度反向计算紧随前向传播的完成而启动，使得缓存的中间结果尽早释放，减少显存消耗。与GPipe类似，在到达同步屏障前，DAPPLE不更新模型参数，而是简单累计求和各个小批次的梯度，所以不会像异步流水线那样需要缓存多个版本的历史模型参数，也不会面临过期参数的问题。DAPPLE的算力利用率理论上与GPipe的相同，但峰值显存消耗平均可节省12%。

1.1.3.3 数据并行

在大数据时代，大规模的训练数据是机器学习训练缓慢的主要矛盾，所以目前在学术界和工业界，主要采用的并行模式是数据并行。在图1-10所示的数据并行框架下，训练数据被切分和存储在多个计算节点，每个计算节点都拥有一份完整的模型副本。计算节点使用本地数据并行训练本地模型，并周期性交换模型数据以实现模型同步。对于不同类型的机器学习算法，模型同步交换的数据内容可能有所区别。例如，对于深度学习等梯度优化类算法，最常见的通信内容是模型参数和梯度；而对于支持向量机和最近邻算法，通信内容则可以是支持向量和关键样本。本书主要讨论基于梯度优化类算法的数据并行方案。

数据并行使用多个计算节点同时处理不同批次的数据，这种并行计算方式能大幅提升分布式机器学习系统的数据吞吐率，有效缩短训练时间。图1-11对比了单机串行和数据并行的训练时间效率。给定三个数据批次，假设每个数据批次都需要四个单位的计算时间（包含前向传播、反向传播和参数更新）。在原始的单机串行模式下，单个GPU独自顺序完成三个数据批次的训练，总计耗费12个单位时间。数据并行模式则将每个数据批次均匀切分给四个GPU，由于计算过程相互无依赖，四个GPU可以并行计算，完成一个数据批次的计算只需要1个单位时间。理想情况下，四卡并行应取得四倍的训练加速，系统可扩展度为1.0。但每次计算完成后，GPU之间需要同步模型参数，假设同步通信延迟为1个单位时间，计算与通信总计需要2个单位时间，加速效果仅达到2倍，此时系统可扩展度为0.5。显而易见，要充分利用GPU的计算资源，必须尽可能提高系统可扩展度。不幸的是，随着集群规模的增大，高网络负载会导致高同步通信延迟，系统可扩展性也会变差。因此，本书将重点介绍分布式机器学习中的通信优化技术，以期取得更优的系统可扩展性。

不同于模型并行和流水线并行需要精心划分模型和放置任务，数据并行可简单适用于任意结构的深度模型。需要注意的是，数据并行更适用于计算密集但参数量小的模型（如卷积神经网络），因为当模型参数量过大时，数据并行的大部分时间都在传输模型参数，高通信代价会引发系统性能瓶颈。尽管如此，我们仍可以通过降低同步频率[51]、压缩传输[52-53]、流量调度[54-55]等手段优化通信效率。

1.1.3.4 混合并行

大规模机器学习训练往往同时面临大数据和大模型的双重挑战。模型并行和流水线并行解决了大模型训练的问题，但面对大数据仍然训练低效。数据并行解决了大数据训练的问题，但当模型过大时，计算节点无法容纳完整的大模型，也限制了数据并行的实用性。因此，在现实的大规模机器学习训练系统中，模型并行、流水线并行、数据并行可能同时存在，这种多并行模式混合的方式称为混合并行。

混合并行的具体实现多种多样，但大致遵循以下规律。（1）参数量大的网络层按神经元切片，如全连接层。（2）具有相似结构的网络块按层切片，如ResNet模型可由不同数量的残差网络块堆叠而成。（3）由于网络块的结构相似，有相近的参数量和计算量，可以结合流水线并行提高算力利用率。（4）当输入不同批次数据时，流水线之间计算与通信相互无依赖，可以数据并行，最后同步模型即可。

图1-12所示为一种混合并行方案。考虑由N个（卷积层，全连接层）基本块堆叠而成的深度网络模型，分布式计算集群总计NM台计算机，每台计算机配备有3个GPU计算卡。首先，依照规律（1）将基本块内的全连接层按神经元切片，由于这种切片方式通信密集，我们将子模型放置单机多卡，如两个全连接层切片放置GPU1和GPU2，卷积层放置GPU0，通过大带宽PCIe通信实现高速多卡模型并行。其次，依照规律（2），我们以基本块为单位切分整个大模型，不同的基本块放置到不同的计算机，并依照规律（3）使用流水线并行串联这些计算机，计算机之间通过数据中心网络通信。最后，以流水线上的N台计算机为一个计算组，整个计算集群可分为M个计算组，这些计算组之间相互无依赖，可以依照规律（4）进行数据并行，进一步提高整体训练效率。

华为和百度飞桨都实现了与图1-12类似的混合并行方案。特别地，由于混合了模型并行、分组参数切片、流水线并行、数据并行四种并行模式，百度飞桨将这种混合方案称为4D混合并行[56]，其对应的分布式训练系统是千亿级语言模型ERNIE的计算基石。其他一些混合方案可参考Mesh-TensorFlow[57]、GNMT[58]、OptCNN[59]、FlexFlow[60]。