DiffusionPipe: Training Large Diffusion Models with Efficient Pipelines

扩散模型已成为图像生成的主要表现者。为了支持训练大型扩散模型，本文研究了扩散模型的管道并行训练，并提出了DiffusionPipe，这是一个同步管道训练系统，提倡创新的管道气泡填充技术，以适应扩散模型的结构特征。目前最先进的扩散模型通常包括可训练的（主干）和不可训练的（例如冻结的输入编码器）部分。我们首先使用动态规划方法统一了单个和多个主干的最佳阶段划分和管道调度，然后提出了一种高效的贪心算法，将不可训练模型部分的计算填充到主干的管道训练的空闲期间，从而实现高训练吞吐量。广泛的实验表明，DiffusionPipe在流行的扩散模型上，可以比管道并行方法提高1.41倍的速度，并比数据并行训练提高1.28倍的速度。

本论文旨在解决大规模扩散模型训练的问题，提出了DiffusionPipe，一个同步管道训练系统，通过创新的管道气泡填充技术，适应扩散模型的结构特征，以支持训练大型扩散模型。

该论文的关键思路是使用动态规划方法统一单个和多个主干的最佳阶段划分和管道调度，然后使用有效的贪心算法将不可训练的模型部分的计算填充到主干的空闲期间，从而实现高训练吞吐量。

论文的实验结果表明，DiffusionPipe可以比管道并行方法提高高达1.41倍的速度，并比数据并行训练提高1.28倍的速度，使用了流行的扩散模型进行测试。此外，该论文提出的创新方法还可以应用于其他大型模型的训练。

最近的相关研究包括使用管道并行训练的其他大型模型的研究，如GPT-3和BERT。