模型并行与数据并行策略

如今训练大型Transformer模型，经常会超出单个加速器（如GPU或TPU）的计算和内存限制。当模型拥有数十亿参数，或数据集需要非常大的批处理尺寸以实现稳定收敛时，将工作负载分散到多个设备上就变得必须。实现此目的的两种主要策略是：数据并行和模型并行。

数据并行

数据并行可能是分布式训练中最直接、最常用的方法。其核心思想简单：在每个可用设备上复制整个模型，并让每个设备处理输入数据批次的不同部分。

工作原理

1. 模型复制：将完整的模型（包含所有参数）复制到每个参与设备上（例如，单台机器上的多个GPU或跨多台机器）。
2. 数据分片：全局训练批次被分割成更小的迷你批次。每个设备接收一个独特的迷你批次。
3. 前向与反向传播：每个设备独立地使用其局部迷你批次执行前向传播以计算损失，然后执行反向传播，根据该局部数据计算模型参数的梯度。
4. 梯度同步：这是重要的通信步骤。在每个设备上计算的梯度会在所有设备之间收集并汇总。一种常见的汇总方法是平均。这通常通过高效的AllReduce通信原语实现。
5. 参数更新：优化器使用汇总的（例如，平均的）梯度来更新模型参数。由于梯度已汇总，所有模型副本都以相同方式更新，从而确保它们保持同步。

数据并行工作流程。模型被复制，数据被分割，梯度在本地计算，跨设备汇总，并用于同步更新所有模型副本。

优点和局限性

数据并行主要优点是其相对简单性，以及随着设备数量增加，训练时间可能实现近乎线性的加速，尤其是在计算量远超通信量时。大多数深度学习框架都提供其实现（例如PyTorch中的torch.nn.parallel.DistributedDataParallel或TensorFlow中的tf.distribute.MirroredStrategy）。

然而，数据并行有一个重要限制：整个模型必须能够放入单个设备的内存中。如果您的Transformer模型对于一个GPU来说太大，那么仅数据并行将无法满足需求。此外，随着设备数量的增加，同步梯度的通信开销可能成为瓶颈，从而降低增加更多设备带来的效益。

模型并行

当模型过大无法放入单个设备时，模型并行就变得必须。我们不是复制模型，而是将模型本身分割到多个设备上。每个设备仅负责存储和计算模型的一部分。

分割模型主要有两种方式：

1. 层间（流水线）并行

这种策略将模型垂直划分。不同的层（或层序列）被分配到不同的设备。数据顺序流经这些设备，形成一个处理流水线。

• 机制：输入数据（或来自前一阶段的激活）进入负责第一组层的设备。这些层的输出激活随后传递给处理下一组层的设备，以此类推，直到产生最终输出。反向传播则沿相反方向流动。
• 挑战 - 流水线气泡：简单的实现会导致设备资源的大量闲置。当一个设备（阶段）正在处理一个批次时，其他设备可能处于空闲状态，等待输入或已完成处理上一个批次。这种空闲时间被称为“流水线气泡”。
• 缓解 - 微批处理：为了减少气泡开销，输入迷你批次被进一步分割成更小的微批次。这些微批次被顺序送入流水线，使得不同阶段能够并发处理不同的微批次，从而重叠计算并减少空闲时间。GPipe或PipeDream等框架为此实现了复杂的调度。

流水线并行将模型层分割到不同设备上。数据顺序流动。在没有微批处理的情况下（简化显示），设备会经历空闲时间（“气泡”）。

2. 层内（张量）并行

这种策略将模型水平划分。它涉及将单个大型层（如自注意力或FFN中的权重矩阵）内部的计算分割到多个设备上。

• 机制：考虑一个大型矩阵乘法 $Y = XA$。并非在单个设备上计算，我们可以将矩阵 $A$ 按列分割（$A = [A_1 | A_2]$），并在设备1上计算 $XA_1$，在设备2上计算 $XA_2$。然后将结果 $Y_1 = XA_1$ 和 $Y_2 = XA_2$ 连接起来形成 $Y = [Y_1 | Y_2]$。类似的分割可以按行应用，或应用于Transformer块内的其他操作。
• 要求：这需要参与计算单个层的设备之间有很大的通信带宽，因为中间激活通常需要交换（例如，使用AllGather或ReduceScatter操作）。
• 应用场景：张量并行对于具有特大型单个层的模型很有效，在这种情况下，仅流水线并行可能不足或导致阶段工作负载不平衡。它通常在流水线的特定阶段使用。NVIDIA的Megatron-LM等库提供了高效的实现。

优点和局限性

模型并行使得训练那些超出单个设备内存容量的模型成为可能。流水线并行通常更容易理解，但受气泡开销影响，需要微批处理以提高效率。张量并行可以处理非常大的层，但要求高设备间带宽并增加了实现复杂性。这两种形式通常比数据并行需要更细致的实现和调试。

混合策略和高级技术

实际中，训练目前先进的大型语言模型通常涉及组合这些策略。一种常见配置是使用流水线并行将层块分布到节点上，并使用张量并行分割每个流水线阶段内的大层。数据并行随后常应用于这种模型并行配置之上，在多组设备上复制整个流水线/张量分割模型以并发处理更多数据。这有时被称为3D并行（数据、流水线、张量）。

此外，诸如ZeRO（零冗余优化器）及其框架实现（例如DeepSpeed、PyTorch FSDP - 完全分片数据并行）提供了精巧的组合。它们作用类似数据并行，但不仅分片数据，还分片优化器状态、梯度，以及可选地将模型参数本身分片到数据并行工作器中。这大幅减少了每个设备的内存占用，使得数据并行能够扩展到比以前大得多的模型，有时甚至能消除中等大小模型对复杂流水线或张量并行的需求。

选择合适的策略

• 从数据并行开始：如果您的模型适合单个设备但训练太慢，数据并行（可能结合ZeRO/FSDP）通常是首先尝试的最简单策略。
• 内存受限时使用流水线并行：如果模型激活内存或参数数量超过单个设备容量，则需要流水线并行。仔细考虑阶段数量和微批次大小，以平衡负载并最大程度减少气泡开销。
• 针对大型层使用张量并行：如果特定层（如大型嵌入表或FFN）即使在流水线阶段内也是瓶颈，请为这些层引入张量并行，确保足够的设备间带宽。
• 组合以实现扩展：对于真正大型的模型（数千亿或数万亿参数），结合数据、流水线和张量并行的混合方法（通常由专门的库管理）通常是必须的。

理解这些分布式训练模式对于有效地使用大型Transformer架构非常重要。尽管深度学习框架提供了实现这些策略的工具，但掌握数据流、通信模式和潜在瓶颈的底层机制，能让您选择正确的方法并优化针对特定模型和硬件配置的训练过程。