流水线并行是一种用于深度神经网络的分布式训练策略，它将模型的层划分为多个顺序阶段，每个阶段分配到不同的设备（如GPU或TPU）。它通过将模型并行与数据并行相结合，以软件流水线的方式处理微批次，旨在提高训练吞吐量并减少朴素逐层并行所带来的空闲时间（即气泡）。
背景与动机
训练大型神经网络，尤其是大型语言模型和其他深度学习模型，需要巨大的计算资源。随着模型规模增长到数十亿甚至数万亿参数，单设备训练因内存和计算限制而变得不可行。数据并行在每个设备上复制整个模型并分割数据，但当模型超出单个设备的内存时便无法使用。模型并行则将模型本身拆分到多个设备上，但朴素的逐层分配会导致严重的资源利用不足：同一时间只有一个设备在计算，其他设备处于等待状态，从而造成低吞吐量和高空闲时间。
流水线并行通过将模型划分为多个阶段（通常是连续的层组）并同时处理多个微批次来解决这一问题。其关键思想是让不同阶段的计算重叠进行：当阶段1在处理微批次2时，阶段2可以处理微批次1，依此类推。这类似于流水线生产线，每个阶段持续处理传入的数据并将结果传递给下一阶段。
核心概念
流水线并行的基本工作单元是微批次。一个大的训练批次被拆分为多个微批次，依次流过流水线的各个阶段。每个阶段对其分配的层执行前向和反向传播。通过依次喂入微批次来“填充”流水线，经过一段预热期后，所有阶段可以同时忙碌，从而实现高吞吐量。
流水线的效率通常用气泡比例来衡量，即阶段空闲的时间占比。对于具有p个阶段和m个微批次的流水线，当m远大于p时，理想吞吐量得以实现，相对气泡开销也会减小。然而，增加微批次数量会增加激活值的内存占用，从而产生权衡。
主要方案
GPipe
Google于2019年提出了GPipe。它将模型划分为多个阶段，并采用简单调度：每个微批次依次完成所有阶段的前向传播，再按相反顺序完成所有阶段的反向传播。GPipe使用同步梯度更新，即所有阶段需等待所有微批次完成后才更新权重。这保证了梯度的一致性，但会引入随阶段数量增长的气泡。GPipe还要求要么重新计算激活值，要么存储它们，从而在计算和内存之间进行权衡。
PipeDream
PipeDream来自Microsoft研究院（2019年），采用异步调度，每个阶段尽可能快地处理微批次，并交替执行前向和反向传播。这减少了气泡，但引入了权重陈旧性问题：不同阶段可能使用不同版本的权重，从而影响收敛性。PipeDream还使用一种称为“权重暂存”的技术，为不同微批次维护多个权重版本，但这会增加内存开销。
PipeDream-2BW及变体
PipeDream-2BW（2020年）通过使用两个权重缓冲区（一个用于前向，一个用于反向）改进了PipeDream，减少了内存占用并改善了收敛性。其变体如V-Pipe和PipeMare则探索了不同的同步策略。
Megatron-LM与交错调度
NVIDIA的Megatron-LM（2019年）将张量并行（拆分单个层）与流水线并行相结合。随后，Megatron-2（2020年）引入了交错调度，将模型划分为比设备数量更多的阶段，使每个设备处理多个阶段。这减少了气泡大小，但增加了通信开销。
Chimera及其他近期方案
Chimera（2021年）采用双向流水线进一步减少气泡。其他方案如PTD-P（2021年）和ZeroBubble（2023年）则探索了更复杂的调度方式，以消除或减少流水线气泡。
数学与实践考量
流水线并行常与其他并行策略结合使用。张量并行将单个层拆分到多个设备，而流水线并行则纵向划分模型。数据并行则将流水线复制到多组设备上。这种混合方法已成为训练大型模型的标准做法，例如Megatron-Turing NLG和OpenAI的GPT-4等系统均采用了此类策略。
流水线深度（阶段数量）和微批次大小的选择会影响性能。更深的流水线可减少每阶段的内存需求，但会增加通信和气泡开销。更大的微批次能提高设备利用率，但会消耗更多内存。最优配置取决于模型大小、设备内存以及设备间互连带宽等因素。
应用与影响
流水线并行已在训练一些最大规模的AI模型中发挥了关键作用。例如，GPT-3（1750亿参数）便是通过模型并行与流水线并行的组合进行训练的。Anthropic的Claude模型和Google的PaLM也采用了类似技术。此外，主流框架均提供了相应支持，包括PyTorch（通过torch.distributed.pipeline模块）、TensorFlow（通过tf.distribute）以及JAX（通过pjit）。
挑战与局限
流水线并行也面临若干挑战。阶段之间传输激活值和梯度会产生通信开销，在互连速度较慢时可能成为瓶颈。各阶段计算时间不均会导致负载失衡，降低整体效率。为反向传播存储激活值会带来内存压力，虽然可通过激活重计算缓解，但这会增加计算成本。此外，异步方案可能因权重陈旧而导致收敛问题，而同步方案则会引入气泡开销。
未来方向
关于流水线并行的研究仍在持续，重点在于减少气泡、降低内存占用以及适应异构硬件。自动阶段划分、动态负载均衡以及与混合专家模型的集成都是活跃的研究方向。随着超大规模模型的兴起以及对云集群高效训练的需求，流水线并行很可能仍将是分布式训练系统中的关键组成部分。