6. 分布式训练

分布式训练主要为了解决单节点设备的算力和内存不足的问题。通常主流的训练框架会提供对应的分布式训练的接口和使用方式。MXMACA软件栈通过对主流训练框架的支持，提供分布式训练的能力。

此外对于大型和超大型模型来说，其本身的模型特点和超出常规的参数量，使得只能通过分布式的算力基础设施来做训练。对此在PyTorch训练框架的基础上，涌现出了类似于Megatron和DeepSpeed这些框架，降低大语言模型训练的门槛。所有的技术底层都离不开分布式训练的各种并行方式和组合策略。本章简要介绍各种并行的方式。

6.1. 数据并行

数据并行的基本使用场景是需要更大的batch来训练某个模型，但是过大的batch会导致显存的OOM，因此单机多卡甚至是多机多卡的数据并行训练将是一个很好的方式。大的batch会被切分到mini-batch供给单卡训练，每个单卡都是一个完整的训练实例，通过对每个step的所有训练实例的权重进行聚合操作并求取均值，然后用这个均值来更新单个训练实例的梯度。

6.2. 模型并行

随着模型规模和参数量的激增，特别是大语言模型的流行，模型并行将是更好的选择。通过将模型进行拆分，放置在不同的训练卡上以满足训练要求。通常有两种类型的并行：张量并行和流水线并行。

6.2.1. 张量并行

张量并行指的是将计算图中层内的参数切分到不同设备，即层内并行，如矩阵-矩阵乘法的拆分。张量并行的关键点在于如何切分模型到不同的设备上和如何保证切分的正确性，因此需要额外的通信开销来保证计算的正确性。例如，对于矩阵乘操作的切分，可以按照左右矩阵的行或列进行切分。不同的切分方式将影响在执行后续操作前是否需要reduce操作以及各种通信和计算的开销。

6.2.2. 流水线并行

流水线并行指的是按模型layer层切分到不同设备，即层间并行。在前向传递过程中，每个设备将中间的激活传递给下一个阶段。在后向传递过程中，每个设备将输入张量的梯度传回给前一个流水线阶段。这允许设备同时进行计算，并增加了训练的吞吐量。流水线并行训练会有一些设备参与计算的冒泡时间，可能导致计算资源的浪费。

6.3. 混合并行

通常单一的并行方式并不能得到最优的资源利用率，特别是在超大模型的训练过程中。为了进一步消除计算设备的冒泡时间，会采取混合的并行方式，即数据并行和模型并行协作的方式来提高训练的效率。

以训练GPT-3 为例，它首先被分为 64 个阶段，进行流水并行。每个阶段都运行在 6 台主机上。在6台主机之间，进行的是数据并行训练；每台主机有 8 张 GPU 显卡，同一台机器上的8张 GPU 显卡之间是进行模型并行训练。