第8章 分布式训练

 

8.1 为什么分布式训练越来越流行

近年来，深度学习已广泛应用于多个领域，如计算机视觉、语言理解、语音识别和广告推荐等。这些领域的共同特点是模型规模越来越大。例如，GPT-3模型的参数量达到了1750亿，使用1024张80 GB的A100显卡，完整训练GPT-3的时间仍然需要约一个月。

随着模型规模的扩大，对硬件（算力、内存）的需求越来越高。然而，受到物理定律和内存墙的限制，单一设备的算力和内存已经无法满足日益增长的需求。为了克服算力不足的问题，人们开始使用多节点集群进行分布式训练，从而提升整体算力，分布式训练变得尤为重要。

8.2 常见的并行策略

仅仅通过简单地堆叠更多机器并不一定能够提高算力。在深度学习中，训练过程不仅仅是将任务分配到多个设备，还涉及设备之间的数据传输。因此，要实现高效的分布式训练，必须协调计算与通信的关系。

为了说明不同的并行策略，我们以矩阵乘法为例进行解释。

8.2.1 数据并行

数据并行是将输入数据 x 切分成多个部分，每个设备上的模型参数 w 保持一致。假设输入数据 x 形状为 4×5，模型参数 w 形状为 5×8，那么矩阵乘法的输出形状为 4×8。

在数据并行中，数据 x 会被按照某一维度（如第0维）切分到多个设备上，每个设备都拥有完整的模型 w。假设将数据 x 切分为两部分，分别传到两个设备，每个设备上的计算会输出形状为 2×8 的部分。最终，这些部分会被拼接，得到完整的输出。

关键问题：

• 数据并行的挑战在于反向传播过程中，每个设备上的梯度（∂loss/∂w）会有所不同。如果不同步这些梯度，设备上的模型会变得不一致。因此，在反向传播过程中需要进行 AllReduce 操作，确保每个设备上的模型始终保持一致。

适用场景：

• 当数据集非常大但模型较小时，数据并行通常较为高效，特别是视觉分类任务（如ResNet50）中，数据并行是一种常见的选择。

8.2.2 模型并行

当模型非常大，单个设备的内存不足以容纳完整的模型时，可以采用模型并行策略。模型并行的核心思想是将模型 w 切分到多个设备上，每个设备仅持有模型的一部分，而数据 x 在所有设备上是共享的。

在模型并行中，每个设备都处理相同的数据 x，但模型参数 w 被分配到不同的设备上。例如，模型 w 按第1维度切分到两个设备上，设备间的计算输出也需要拼接以得到完整的结果。

优点：

• 模型并行避免了多个设备之间的梯度同步（AllReduce）的高通信成本。

缺点：

• 每个设备仍然需要接收到完整的数据输入，因此数据会在多个设备间进行广播，这样会引入通信开销。

适用场景：

• 在如BERT等大规模语言模型中，模型并行是常见的策略。

8.2.3 流水并行

当神经网络非常庞大时，除了采用模型并行外，还可以选择流水并行。流水并行将网络切分成多个阶段，每个阶段被分配到不同的计算设备上。各个设备之间通过“接力”方式完成训练，即一个设备的输出会作为下一个设备的输入。

例如，一个逻辑上的4层神经网络（T1至T4）被切分到两个设备上，其中GPU0负责计算T1和T2，GPU1负责计算T3和T4。当GPU0计算完T1和T2后，输出被传递到GPU1，继续完成T3和T4的计算。

优点：

• 流水并行可以有效地将计算任务拆分到多个设备，适合超大规模模型。

缺点：

• 由于涉及设备间的连续传递，可能会引起延迟和带宽问题。

8.2.4 混合并行

混合并行策略结合了多种并行策略的优势，通常用于非常大的模型。例如，GPT-3的训练采用了混合并行策略：

1. 流水并行：将网络划分为64个阶段，每个阶段由6台DGX-A100主机处理。
2. 数据并行：在每台主机之间采用数据并行训练。
3. 模型并行：每台主机的8张GPU显卡之间进行模型并行训练。

这种多层次的并行方案使得大规模模型的训练能够高效进行。

总结

不同的并行策略适用于不同的场景，选择合适的并行策略是提升训练效率的关键。数据并行适用于较小模型和大规模数据集，模型并行适合大模型但较小的数据集，流水并行适合非常庞大的网络，而混合并行则结合了多种策略以满足复杂需求。选择何种策略将依赖于硬件资源、模型大小以及数据分布等多方面的因素。