PyTorch 流水线并行模式设计分析

流水线并行(Pipeline Parallelism)最早是 Google 在 Gpipe 论文中提出的,这种并行训练模式能够充分利用多 GPU 的资源高效地训练评估大模型。目前 PyTorch 最新版本是 2.2,流水线并行的功能是基于 torchgpipe 论文中的设计来实现的,该功能当前还处于试验阶段。 问题背景 大模型无法直接放到单个 GPU 中进行训练,通过模型并行(Model Parallelism)的方法可以把模型进行分片,每一个分片放置到一个 GPU 上,这样能够很好实现模型并行且利用多 GPU 的资源。虽然使用这种较为初级的方式能够实现大模型的训练,但在训练的过程中并不能充分利用 GPU 资源,因为对顺序(Sequential)模型来说它每次只能激活一个 GPU 来进行训练,其它的 GPU 此时是闲置的,所以在底层设备上其实仍然是顺序执行。 例如,对一个有 4 层的顺序(Sequential)神经网络模型,经过模型分片后,训练过程中每一层(或 Subnetwork)放在一个 GPU 上,先进行前向传播计算得到 Loss,然后反向传播计算梯度,如下图所示: 使用这种方式利用 GPU 训练,我们可以看到在训练过程中 GPU 完全没有被充