大语言模型的工程技巧（三）——分布式计算

相关说明

这篇文章的大部分内容参考自我的新书《解构大语言模型：从线性回归到通用人工智能》，欢迎有兴趣的读者多多支持。

本文将讨论如何利用多台机器进行神经网络的分布式训练。利用多台机器来加速大语言模型的训练，是其获得成功的重要原因。

关于其他的工程技巧可以参考：

• 大语言模型的工程技巧（一）——GPU计算
• 大语言模型的工程技巧（二）——混合精度训练

关于大语言模型的内容，推荐参考这个专栏。

一、概述

本文将讨论如何巧妙地借助多台机器来优化模型训练和应用速度。在神经网络领域，常常利用GPU进行模型计算，以迅速提高计算效率。然而，正如大语言模型的工程技巧（一）——GPU计算所述，即使在同一台机器上，跨GPU的数据也无法直接运算。因此，对于分布式运算，多台机器之间的协作机制相当于不同GPU之间（不管它们是否在同一台机器上）的协作机制。为了表述简单，本节后续的讨论都只针对在多个GPU之间的分布式计算。

二、两种并行

模型计算的基础是计算图，因此，模型的分布式计算实质上就是在计算图层面进行分布式运算。关于这一主题，业界出现了两种截然不同的分布式计算方法，分别是数据并行（Data Parallelism）和模型并行（Model Parallelism）。数据并行，也就是梯度累积[TODO]，它根据数据将计算图纵向切分，从而进行并行计算。与之不同，模型并行是将计算图的不同层放置在不同的GPU上进行计算。这可以被形象地理解为：数据并行将计算图从竖直方向切分，而模型并行从水平方向切分，如图1所示。

图1

三、数据并行

在传统的观念里，模型的分布式计算意味着对数据的并行处理。这种方法的核心思想遵循著名的Map/Reduce框架¹模式，如图2所示。首先，数据被智能地分发到各个GPU上。接着，完整的模型被逐一复制到每个GPU上。然后，这些GPU利用各自的数据进行向前传播和反向传播，这一系列步骤类似于“映射”（Map）操作。随后，执行“归约”（Reduce）操作（更确切地说是“All Reduce”操作²）。在这一阶段，算法将每个GPU上的反向传播梯度传递给其他GPU。简而言之，每个GPU都积累了所有GPU计算得出的梯度信息，能够独立地累加梯度，并进行后续的参数迭代更新。由于每个GPU上累加的梯度相同，因此在参数更新后得到的模型也是相同的。持续循环，直到得到最终的模型。这个过程确保了模型的并行训练和参数同步。

图2

从每个GPU的角度来看，尽管每次迭代只处理批次数据中的一部分，但在Reduce阶段，通过梯度的传递，参与模型参数更新的梯度却基于整个批次的所有数据。换句话说，这个阶段汲取了批次中全部数据的智慧。这就好比一份试卷，一个班级的学生各自分工做不同的试题，然后相互交流答案，这样每个学生只解答了部分问题，却获得了全部答案。因此，即使硬件未经升级，GPU的学习速度也会更快，从而加速整个模型的训练过程。借助这种巧妙的分布式计算方式，我们能够汇聚个体的努力，更迅速地训练模型。

四、模型并行

近年来，随着模型规模的持续扩大，针对单个数据的模型计算量变得异常庞大，有时甚至超越了单个GPU的处理能力，导致计算难以进行。为了应对这一挑战，业界开始探索一种全新的分布式计算思路，即模型并行。如图3所示，将计算图的不同层分散到不同的GPU上，以神经网络为例，可以将神经网络的各层分配给不同的GPU。这样，每个GPU只需要负责模型的一部分，只有按照正确的顺序将它们串联在一起，才能构建出完整的模型。在计算过程中，前一个GPU的计算结果将成为后一个GPU的计算图输入，多个GPU合作完成一次计算图的计算。通过多个GPU的协同合作，我们能够有效地处理单个GPU难以胜任的大规模模型的计算。

模型并行不仅可以应对庞大的模型规模带来的挑战，还能够提升模型计算的速度。为了理解这一点，可以将模型并行的过程类比为流水线，GPU是流水线上的一环。如图3所示，在GPU:1处理第一份数据的同时，GPU:0已经开始处理第二份数据了。通过充分利用流水线的并行原理，整个模型的计算速度得到了显著提升。

图3

五、代码实现

上述两种方法并非互斥的选择，而是可以将两者结合使用，以提升计算效率。例如，在数据并行的大框架下，当一台拥有多个GPU的机器对相应数据进行计算时，可以采用模型并行的策略将模型分散到不同的GPU上，从而进一步提升计算速度。

分布式计算本身相当复杂，除了涉及算法层面的代码实现，还涉及集群层面的构建和维护工作，如机器间的通信和错误恢复等。在这两个方面，PyTorch提供了出色的支持。在代码方面，PyTorch提供了3个优秀的封装工具³，分别是torch.distributed、torch.multiprocessing和torch.nn.parallel.DistributedDataParallel，可以帮助我们快速搭建分布式模型，具体的代码实现可以参考这个链接。在集群搭建⁴方面，PyTorch提供了torchrun工具，致力于更轻松地配置集群环境。

---

1. Map/Reduce框架是一种经典的分布式计算模式，整个计算过程分为两个关键阶段：Map和Reduce。它最初由Google提出，并在处理海量数据时取得了巨大成功。这个框架的设计思想旨在将复杂的任务分解成多个简单的子任务，分布在多台机器上并行执行（Map阶段），然后将结果合并（Reduce阶段）以得到最终的计算结果。 ↩︎

2. 在经典的Map/Reduce框架中，Reduce操作只在选定的一台机器上进行，并非在全部机器上执行，因此这里的步骤被称为All Reduce。 ↩︎

3. 这里涉及的3个工具都用于数据并行的情况，若要实现模型并行，则需要自行编写代码。幸运的是，具体的实现并不复杂，所涉及的核心流程是GPU计算中的数据复制。 ↩︎

4. 对于用于机器学习的专用集群（通常为GPU集群），有一些更专业的工具可用于集群的搭建和管理，比如NVIDIA Bright Cluster Manager、Slurm等。这些工具旨在优化集群的性能，确保计算资源得到最大限度的利用。 ↩︎