流水线并行是一种面向深度神经网络的分布式训练策略，它将模型的层划分为多个顺序阶段，每个阶段分配给不同的设备（如GPU或TPU）。它通过以软件流水线方式处理微批次在各阶段间的流动，将模型并行与数据并行相结合，旨在提高训练吞吐量并减少朴素逐层并行所带来的空闲时间（气泡）。
背景与动机
训练大型神经网络，尤其是大型语言模型和其他深度学习模型，需要巨大的计算资源。随着模型规模增长到数十亿或数万亿参数，单设备训练因内存和计算限制而变得不可行。数据并行在每个设备上复制整个模型并划分数据，但当模型超出单个设备内存时便会失效。模型并行则将模型本身拆分到多个设备上，但朴素的逐层放置会导致严重的资源利用不足：每次只有一个设备计算，其余设备等待，导致吞吐量低且空闲时间长。
流水线并行通过将模型划分为多个阶段（通常是连续的层组）并同时处理多个微批次来解决这一问题。其关键洞察在于跨阶段重叠计算：当阶段1处理微批次2时，阶段2可以处理微批次1，依此类推。这类似于装配线，每个阶段持续处理传入的数据并将结果传递给下一阶段。
核心概念
流水线并行的基本工作单元是微批次。一个大的训练批次被拆分成较小的微批次，顺序流经流水线各阶段。每个阶段对其分配的层执行前向和反向传播。流水线通过逐个输入微批次来“填充”，在初始预热期后，所有阶段可以同时忙碌，从而实现高吞吐量。
流水线的效率通常通过气泡比率来衡量，即阶段空闲的时间比例。对于具有p个阶段和m个微批次的流水线，当m远大于p时，可实现理想吞吐量，从而最小化相对气泡开销。然而，增加m也会增加激活的内存使用，形成权衡。
主要方案
已提出多种流水线并行方案，它们在前向和反向传播的调度方式及梯度更新处理上有所不同。
GPipe
Google于2019年推出了GPipe。它将模型划分为多个阶段，并使用简单调度：每个微批次先通过所有阶段进行前向传播，再按相反顺序进行反向传播。GPipe使用同步梯度更新，即所有阶段等待所有微批次完成后再更新权重。这保证了梯度的一致性，但引入了随阶段数量增长的气泡。GPipe还需要重计算激活或存储它们，从而在计算与内存之间进行权衡。
PipeDream
来自微软研究院的PipeDream（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）。
挑战与局限
流水线并行引入了若干挑战。通信开销来自阶段之间传递激活和梯度，在慢速互连上可能成为瓶颈。负载不均衡发生在各阶段计算时间不等时，会降低效率。存储激活用于反向传播所带来的内存压力可以通过激活重计算来缓解，但这会增加计算成本。此外，异步方案可能因权重陈旧而出现收敛问题，而同步方案则会产生气泡。
未来方向
关于改进流水线并行的研究仍在继续，重点关注减少气泡、最小化内存以及适应异构硬件。自动阶段划分、动态负载均衡以及与混合专家模型的集成等技术是活跃领域。大型模型的兴起以及云集群上高效训练的需求，很可能使流水线并行继续成为分布式训练系统的关键组成部分。