PipeFusion: Patch-level Pipeline Parallelism for Diffusion Transformers Inference

作者/机构: Jiarui Fang (ByteDance), Jinzhe Pan (HUST), Aoyu Li (ByteDance), Xibo Sun (Tencent), Jiannan Wang (The University of Hong Kong)

A1 主要贡献

本文提出了一种名为 PipeFusion 的创新并行方法,旨在解决使用扩散变换器(Diffusion Transformers, DiTs)模型生成高分辨率图像时遇到的高延迟问题。

核心问题:随着AI生成内容(AIGC)的发展,DiT模型正逐渐取代U-Net架构。然而,由于注意力计算的二次方延迟增长,DiT在处理长序列(如高分辨率图像)时面临严重的性能瓶颈。现有的并行方法,如张量并行(TP)和序列并行(SP),在应用于DiT时效率不高,前者因激活值体积大而导致通信成本过高,后者则需要在每个注意力层进行数据交换。DistriFusion虽利用了扩散过程中的时间冗余来重叠通信,但应用于DiT时内存效率低下。

研究目标:开发一种新的并行策略,以更有效地利用DiT推理过程中的输入时间冗余,显著降低通信成本和内存消耗,从而加速高分辨率图像的生成。

创新点 (PipeFusion)
1. 分块级流水线并行策略:PipeFusion将图像分割成多个补丁(patch),并将模型的网络层划分到多个GPU上。它采用一种补丁级的流水线并行策略,高效地协调通信与计算。
2. 利用时间冗余减少通信:通过利用连续扩散步骤之间输入的高度相似性,PipeFusion重用前一步的“陈旧”特征图(stale feature maps)为当前流水线步骤提供上下文。这与现有的DiT并行方法(如张量并行、序列并行和DistriFusion)相比,极大地减少了通信开销,因为它仅需在每个设备上各阶段的边界传输输入和输出,而非每层都传输激活值。
3. 高内存效率:通过将模型参数分布在不同设备上,PipeFusion提升了内存效率。与DistriFusion需要每个设备缓存完整的K和V激活不同,PipeFusion所需的注意力K、V缓存大小减少为1/N(N为并行度),使其非常适合像12B参数的Flux.1这样的大型DiT模型。
4. 更高的生成精度:相比于DistriFusion,PipeFusion在更长的时间内使用更新的(fresh)K和V值,从而实现了更高的图像生成精度。

实验结果表明,在8块L40 PCIe GPU上,PipeFusion在Pixart、Stable-Diffusion 3和Flux.1模型上均取得了当前最先进的性能。

A3 背景知识

扩散模型与扩散变换器 (DiTs):扩散模型使用一个噪声预测深度神经网络(DNN),表示为$\epsilon_{\theta}$,来生成高质量图像。该过程从纯高斯噪声$x_T \sim N (0, I)$开始,通过多个步骤迭代去噪,最终生成图像$x_0$,其中$T$是总扩散步数。如图1所示,在每一步$t$,模型$\epsilon_{\theta}$接收带噪图像$x_t$、时间步$t$以及一个附加条件$c$(如文本、图像)作为输入,以预测$x_t$中的噪声$\epsilon_t$。前一步的图像$x_{t-1}$通过以下方式获得:

$$x_{t-1} = \text{Update}(x_t, t, \epsilon_t), \quad \epsilon_t = \epsilon_\theta(x_t, t, c).$$

函数Update的具体形式取决于采样器(例如,DDIM【索引编号17,Denoising diffusion implicit models,2020,arXiv】或DPM【索引编号18,Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps,2022,Advances in Neural Information Processing Systems】),并且涉及逐元素操作。最终,$x_0$通过变分自编码器(VAE)【索引编号19,Auto-encoding variational bayes,2013,arXiv】从潜在空间解码到像素空间。推理延迟的主要贡献者是通过$\epsilon_{\theta}$的前向传播。$\epsilon_{\theta}$的架构正在从U-Net【索引编号7,U-net: Convolutional networks for biomedical image segmentation,2015,Medical image computing and computer-assisted intervention–MICCAI 2015】转向扩散变换器(DiTs)【索引编号20,Megascale: scaling large language model training to more than 10,000 gpus,2024,arXiv;索引编号21,Pixart-α: Fast training of diffusion transformer for photorealistic text-to-image synthesis,2023,arXiv;索引编号22,Latte: Latent diffusion transformer for video generation,2024,arXiv】,这是由显示模型参数和训练数据增加能带来性能提升的缩放定律驱动的。与使用卷积层捕捉空间层次结构的U-Nets不同,DiTs将输入分割成潜在的补丁,并使用变换器的自注意力机制【索引编号23,Attention is all you need,2017,Advances in neural information processing systems】来建模这些补丁内部和之间的关系。输入的带噪潜在表示被分解为补丁,嵌入为令牌,然后送入一系列DiT块,这些块通常包括多头自注意力、层归一化和前馈网络。

图1: DiT推理的工作流程。
图1: DiT推理的工作流程。

输入时间冗余:扩散模型需要在输入图像或视频中迭代预测噪声。最近的研究强调了输入时间冗余的概念,这是指在连续的扩散时间步中观察到的输入和激活值的相似性【索引编号11,Frdiff: Feature reuse for exquisite zero-shot acceleration of diffusion models,2023,CoRR;索引编号12,Deepcache: Accelerating diffusion models for free,2024,Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition】。最近的一项研究【索引编号13,Learning-to-cache: Accelerating diffusion transformer via layer caching,2024,arXiv】进一步调查了这种相似性在不同层和时间步的分布。基于这种冗余,一系列研究通过缓存激活值并在后续扩散时间步中重用它们来减少计算。例如,在U-Net架构的背景下,DeepCache更新低级特征,同时重用缓存中的高级特征【索引编号12,Deepcache: Accelerating diffusion models for free,2024,Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition】。此外,TGATE在交叉注意力沿扩散过程收敛后缓存其输出【索引编号12,Deepcache: Accelerating diffusion models for free,2024,Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition】。相比之下,在DiT模型领域,∆-DiT提出在早期采样阶段缓存后部的DiT块,在后期阶段缓存前部的DiT块【索引编号24,\delta-dit: A training-free acceleration method tailored for diffusion transformers,2024,arXiv】。PAB【索引编号15,Real-time video generation with pyramid attention broadcast,2024,arXiv】利用U形注意力模式,通过金字塔式广播方法来减轻时间冗余。最后,DiTFastAttn【索引编号14,Ditfastattn: Attention compression for diffusion transformer models,2024,arXiv】识别了三种类型的冗余,如空间冗余、时间冗余和条件冗余,并提出了一种注意力压缩方法来加速生成。

A2 方法细节

本节中,我们深入探讨了张量并行(TP)、序列并行(SP)以及一种利用输入时间冗余的异步SP方法,即DistriFusion,在DiT中进行并行推理的应用。随后,我们介绍了一种名为PipeFusion的新型分块级流水线并行方法,该方法相比现有技术显著提高了通信效率。

在本文的其余部分,使用以下符号:$p$表示序列长度,对应于潜在空间中的像素数量。$h_s$表示模型的隐藏大小。$L$代表网络层数。$N$是计算设备的数量。

3.1 张量并行与序列并行

张量并行 (TP):张量并行(TP)按列和按行对DiT的线性层进行分区。这种方法每层需要两次all-reduce操作,每层产生$4O(p \times h_s)$的通信开销(见表1)。通信成本随序列长度线性扩展,导致可扩展性差,特别是对于像0.6B参数的PixArt模型这样的大型模型(图8)。此外,TP对非标准架构(如Flux/SD3中的MM-DiT)的实现带来了挑战。

序列并行 (SP):序列并行(SP)将图像划分为多个补丁,每个设备处理一个补丁作为输入。由于注意力机制需要全局交互,设备之间必须通信部分的查询(Q)、键(K)和值(V)。如图2所示,两种著名的SP实现是DeepSpeed-Ulysses【索引编号9,Deepspeed ulysses: System optimizations for enabling training of extreme long sequence transformer models,2023,arXiv】(SP-Ulysses)和Ring-Attention【索引编号10,Ring attention with blockwise transformers for near-infinite context,2023,arXiv】(SP-Ring)。SP-Ulysses在并行注意力计算期间使用All2All通信在序列和隐藏大小分区之间切换,而SP-Ring通过点对点的K/V子块传输实现分布式Flash Attention。值得注意的是,SP-Ulysses和SP-Ring可以结合形成USP(统一序列并行)【索引编号25,A unified sequence parallelism approach for long context generative ai,2024,arXiv】,这通常比单独使用任一方法更有效。

图2: DistriFusion与序列并行方法(DeepSpeed-Ulysses和Ring-Attention)在注意力层上的比较。
图2: DistriFusion与序列并行方法(DeepSpeed-Ulysses和Ring-Attention)在注意力层上的比较。

3.2 DistriFusion:异步序列并行

DistriFusion:异步序列并行:DistriFusion是一种异步序列并行方法,它利用了输入数据中的时间冗余。它使用来自前一时间步的稍微过时(stale)的激活值,而不是在每一步都要求最新的激活值。尽管如此,DistriFusion仍能保持图像生成精度,且质量上没有任何可察觉的损失。

DistriFusion的DiT扩展:虽然DistriFusion最初应用于基于U-Net的SDXL模型【索引编号26,Sdxl: Improving latent diffusion models for high-resolution image synthesis,2023,arXiv】,但它也可以扩展到DiT模型。如图2所示,DistriFusion也是通过沿序列维度进行分区来操作的。然而,关键区别在于它使用异步的all-gather操作从远程设备收集K和V的激活值。由于这种通信的异步性,当前扩散步骤的激活值不能立即可用;相反,它们在下一个扩散步骤期间才变得可用。

DistriFusion的注意力计算与通信重叠:DistriFusion利用这一点,将在时间步$T+1$的K和V激活值中的$N-1/N$部分,与在扩散时间步$T$的本地K和V的$1/N$部分结合起来。这个组合被用来与扩散时间步$T$的本地查询进行注意力操作计算。时间步$T$的K和V的通信被有意地与时间步$T$的网络前向计算重叠,从而在计算期间隐藏了通信开销。

DistriFusion的内存开销:然而,DistriFusion以增加内存使用为代价实现了这种通信-计算重叠。每个计算设备都需要维护通信缓冲区,以存储K和V激活的完整空间形状,总计为$A_L$。因此,DistriFusion的内存成本不会随着计算设备的增加而减少。

3.3 PipeFusion: 分块级流水线并行

PipeFusion:分块级流水线并行:PipeFusion同时对输入序列和DiT主干网络的层进行分区,如图3顶部所示。PipeFusion沿着数据流的方向对DiT模型进行分区。因此,每个分区包含一组连续的层,并部署在一个GPU上。PipeFusion还将输入图像划分为M个不重叠的补丁,以便每个GPU并行处理一个补丁及其分配的层级的计算,并且扩散例程以流水线方式高效运行。流水线操作需要在设备之间进行同步,如果DiT的工作负载在每个GPU上均匀分布,这将是高效的。由于DiT包含许多相同的变换器块,因此很容易实现均匀分区。

图3: 上图:输入和DiT主干网络的划分策略。下图:PipeFusion作为分块级流水线并行的工作流程。
图3: 上图:输入和DiT主干网络的划分策略。下图:PipeFusion作为分块级流水线并行的工作流程。

PipeFusion流水线工作流程:假设我们当前处于扩散时间步$T$,前一个时间步是$T+1$,因为扩散过程是沿着时间步的逆序进行的。图3中的例子展示了$N=4$和$M=4$时的流水线工作流程,其中时间步$T$的补丁激活值被高亮显示。由于输入的时间冗余,一个设备不需要等待接收当前时间步$T$的完整空间形状激活值来开始其自身阶段的计算。相反,它使用来自前一个时间步的陈旧激活值为当前计算提供上下文。例如,在微步5中,补丁P0和P1的激活值来自时间步$T$,而其他的来自时间步$T+1$。通过这种方式,每个设备都维护完整的激活值,并且在流水线初始化后没有等待时间。气泡(bubbles)只存在于流水线的开始阶段。假设总扩散时间步数为$S$,则流水线的有效计算比率为 $M \cdot S / (M \cdot S + N - 1)$。由于$S$通常为了高质量图像生成而被选为一个较大的数,因此有效计算比率很高。例如,当$M=N=4$且$S=50$时,该比率等于98.5%。只要补丁数$M$大于并行度$N$,除了初始启动开销外,就不会产生气泡。我们建议设置$M=N$。我们已在附录B节中分析了$M$对最终性能的影响。

通信与计算重叠:此外,在PipeFusion中,一个设备通过异步P2P将微步补丁的激活值发送给后续设备,从而实现了通信与计算的重叠。例如,在微步4,设备0接收来自扩散时间步$T$的P1的激活值,同时计算时间步$T$的P0的激活值。同样,在下一个微步,P0到设备1的传输可以被设备上P1的计算所隐藏。

预热阶段 (Warmup):与DistriFusion类似,在执行流水线之前,我们通常会同步进行几次扩散迭代,称为预热步骤。在预热阶段,补丁是顺序处理的,导致效率较低。尽管预热步骤不能以流水线方式执行,但其工作流程相对于整个扩散过程较小,因此对性能的影响可以忽略不计。我们在附录B节中分析了预热的影响并提出了相应的解决方案。

生成质量优势:从理论上讲,考虑到新鲜激活区域的面积,PipeFusion在生成质量方面优于DistriFusion。如图4所示,在单个扩散时间步内,随着流水线微步从扩散时间步4推进到8,PipeFusion持续增加新鲜激活的区域。相比之下,在整个扩散过程中,DistriFusion在总共M个补丁中始终只维持一个补丁的新鲜区域。因此,PipeFusion的结果比DistriFusion更接近原始扩散方法。

图4: 图3中扩散时间步T期间激活的新鲜部分。深灰色代表新鲜数据,浅灰色代表陈旧数据。
图4: 图3中扩散时间步T期间激活的新鲜部分。深灰色代表新鲜数据,浅灰色代表陈旧数据。

预热阶段对流水线的影响:输入时间冗余的利用需要一个短暂的预热阶段,因为扩散合成的行为在去噪过程中会发生质的变化。因此,DistriFusion在切换到其异步模式之前,会使用标准的同步补丁并行执行几个预热步骤。对于PipeFusion,预热步骤不能流水线化,因此会引入气泡。然而,由于预热步骤的数量通常远小于总扩散迭代次数,它们对整体效率的影响是有限的。我们在B节中分析了这种影响。

与现有流水线并行的区别:PipeFusion代表了与以往流水线并行化方法(如GPipe【索引编号27,Gpipe: Efficient training of giant neural networks using pipeline parallelism,2019,Advances in neural information processing systems】和TeraPipe【索引编号28,Terapipe: Token-level pipeline parallelism for training large-scale language models,2021,International Conference on Machine Learning】)的根本不同。与GPipe沿批次维度分割并依赖于多请求高并发不同,PipeFusion沿序列维度分割,为单个图像提供高效的流水线执行。与专为语言模型中的因果注意力设计的TeraPipe不同,PipeFusion专为DiT的全注意力机制量身定制,并利用时间冗余来同时提高效率和保持生成质量。

3.4 不同并行方法之间的比较

PipeFusion的通信与内存效率优势:与现有的并行解决方案相比,PipeFusion在通信和内存使用方面表现出卓越的效率。首先,PipeFusion在计算设备之间仅传输作为一系列连续Transformer层(属于一个阶段)的输入和输出的激活值,与其它需要为所有L层传输KV激活的方法相比,显著降低了通信带宽需求。具体来说,PipeFusion的总通信成本为$2O(p \times h_s)$,与层数$L$无关。其次,PipeFusion中的每个设备仅存储其特定层的$1/N$参数。由于使用陈旧的KV进行注意力计算需要每个设备为其上的$L/N$层维护完整的空间KV,其开销显著小于DistriFusion,并且随着设备数量的增加而减少。

表1: 不同DiT并行方法在单个扩散时间步上的比较。Overlap表示通信与计算之间的重叠。
表1: 不同DiT并行方法在单个扩散时间步上的比较。Overlap表示通信与计算之间的重叠。

理论成本总结:接下来,我们在表1中总结了现有方法的理论成本。通信成本是通过传输元素数量与与通信类型相关的算法带宽(algobw)因子3的乘积计算得出的。对于AllReduce、AllGather和AllToAll等集体算法,相应的algobw因子分别为$2(n-1)/n$、$(n-1)/n$和1。在表中,为简化起见,我们将$O((n-1)/n)$项近似为$O(1)$。

各方法通信成本比较:在所有方法中,只要$N < 2L$,PipeFusion的通信成本最低,这很容易满足,因为网络层数$L$通常很大,例如在StableDiffusion-3中$L = 38$。此外,PipeFusion可以重叠通信和计算。Ulysses序列并行的通信成本随着$N$的增加而降低,优于其余三种方法。然而,它的通信无法被计算隐藏。Ring序列并行和DistriFusion具有相似的通信成本和重叠行为。区别在于重叠的范围:Ring序列并行中的计算在注意力模块内部重叠,而DistriFusion中的计算则在整个前向传播过程中重叠。

各方法内存成本分析:为了分析内存成本,我们用$P$表示模型参数的总数。在PipeFusion和张量并行中,内存成本随着使用更多GPU而降低,这对于快速增长的DiT模型大小是一个很好的特性。PipeFusion和DistriFusion都为每个Transformer层维护一个KV缓冲区,导致显著的激活内存开销,特别是对于长序列。PipeFusion中的KV缓冲区随着设备数$N$的增加而减小,而DistriFusion则没有这种减少。

A4 实验环境

硬件配置
* 主要平台:一个由8块L40-48GB GPU组成的集群,通过PCIe Gen4x16连接。
* 对比平台(附录):8块A100 GPU,通过NVLink(高带宽)连接。

模型架构
* Pixart 【索引编号21, Junsong Chen et al., Pixart-α: Fast training of diffusion transformer for photorealistic text-to-image synthesis, 2023, arXiv; 索引编号3, Junsong Chen et al., Pixart-σ: Weak-to-strong training of diffusion transformer for 4k text-to-image generation, 2024, arXiv】:主干是一个包含6亿(0.6B)参数的Transformer模型,基于原始DiT框架,并集成了用于合并文本条件的交叉注意力模块。
* Stable-Diffusion-3 (SD3-medium) 【索引编号2, Patrick Esser et al., Scaling rectified flow transformers for high-resolution image synthesis, 2024, Forty-first International Conference on Machine Learning】:利用多模态(MM)-DiT,其主干Transformer模型包含20亿(2B)参数。使用20步的FlowMatchEulerDiscreteScheduler。
* Flux.1-dev 【索引编号1, BlackForestLabs, Announcing black forest labs, 2024, https://blackforestlabs.ai/ announcing-black-forest-labs/】:主干Transformer模型包含120亿(12B)参数,采用增强的MM-DiT设计。使用28步的FlowMatchEulerDiscreteScheduler。

软件配置
* PyTorch 2.4.1
* CUDA Runtime 12.1.105
* diffusers 0.30.3

数据集
* COCO Captions 2014 【索引编号37, Xinlei Chen et al., Microsoft coco captions: Data collection and evaluation server, 2015, arXiv】:用于评估Fréchet Inception Distance (FID) 分数。从验证集中采样30,000张图像作为参考,并为每个实验生成30,000张图像进行比较。

A4 实验结果

本节比较了不同并行方法在8块通过PCIe Gen4连接的L40 GPU上的性能。所有延迟数据均不包括VAE解码部分,且为5次运行的平均值。对于DistriFusion和PipeFusion,默认使用1个预热步骤。

4.1.1 延迟评估

  • Pixart (图5):在1024px、2048px和4096px分辨率下,PipeFusion均表现出最低的延迟,分别比次优方法快1.48倍、1.55倍和1.16倍。SP-Ulysses因PCIe带宽瓶颈延迟最高。USP(SP-Ulysses和SP-Ring的混合)优于单独的SP方法。DistriFusion在4096px任务上出现内存溢出(OOM)。结合CFG并行后,PipeFusion(cfg=2, pipefusion=4)与单GPU基线相比,实现了高达8.59倍的延迟降低。
图5: 在20步DPMSolverMultistepScheduler下,各种并行方法在两个图像生成任务上对Pixart的延迟。
图5: 在20步DPMSolverMultistepScheduler下,各种并行方法在两个图像生成任务上对Pixart的延迟。
  • Stable Diffusion 3 (图6):在1024px和2048px分辨率下,PipeFusion优于最佳的序列并行方法(USP),分别取得了1.57倍和1.30倍的加速。结合CFG并行,与单GPU基线相比,实现了高达8.16倍的加速。由于MM-DiT架构的复杂性,未实现张量并行。
图6: 在20步FlowMatchEulerDiscrete调度器下,各种并行方法在两个图像生成任务上对SD3-medium的延迟。
图6: 在20步FlowMatchEulerDiscrete调度器下,各种并行方法在两个图像生成任务上对SD3-medium的延迟。
  • Flux.1-dev (图7):在1024px和2048px分辨率下,PipeFusion同样优于最佳的序列并行方法(USP),分别取得了1.23倍和1.25倍的加速。与单GPU基线相比,实现了高达5.79倍的加速。
图7: 在28步FlowMatchEulerDiscrete调度器下,各种并行方法在三个图像生成任务上对Flux.1-dev的延迟。
图7: 在28步FlowMatchEulerDiscrete调度器下,各种并行方法在三个图像生成任务上对Flux.1-dev的延迟。
  • 可扩展性 (图8):通过将GPU数量从1扩展到8,评估了三种模型的可扩展性。PipeFusion在所有五个场景中始终表现出最佳的可扩展性。即使在计算密集型的4096px任务上,PipeFusion的通信开销占比仅为4.6%,远低于SP的17.9%,绝对通信时间减少了78%,证明了其在抑制通信开销方面的鲁棒性。
图8: 三种DiT在8×L40上的可扩展性分析:最左侧图表展示了Pixart在1024px、2048px和4096px图像上的性能。右二图表显示了SD3在1024px上的性能。最右侧图表说明了Flux.1-dev在1024px上的性能。图中标注了PipeFusion曲线相对于1 GPU基线的加速比。
图8: 三种DiT在8×L40上的可扩展性分析:最左侧图表展示了Pixart在1024px、2048px和4096px图像上的性能。右二图表显示了SD3在1024px上的性能。最右侧图表说明了Flux.1-dev在1024px上的性能。图中标注了PipeFusion曲线相对于1 GPU基线的加速比。

4.1.2 内存效率

  • Pixart (图9a):DistriFusion的内存消耗随图像分辨率增加而增长。相比之下,PipeFusion的内存占用与张量并行和序列并行相当,并且在8192px分辨率下内存消耗最低。

  • SD3-medium 和 Flux.1-dev (图9b):PipeFusion的内存成本显著低于序列并行。对于12B参数的Flux.1-dev模型,PipeFusion的总内存使用量仅为序列并行的32%(1024px)和36%(2048px),尽管其KV缓冲区增加了激活值的内存占用,但与模型参数的内存节省相比,这一增长相对较小。

图9: 不同模型和分辨率下GPU内存使用情况的比较。
图9: 不同模型和分辨率下GPU内存使用情况的比较。

4.2 质量结果

  • FID分数 (表2):Fréchet Inception Distance (FID) 分数评估结果显示,PipeFusion在相同设备配置下始终优于DistriFusion,取得了更低的FID分数(分数越低表示质量越好)。这与理论分析(第3.3节)相符,即PipeFusion利用了更新鲜的激活值。
表2: Pixart和Flux.1上并行方法的FID分数。w/ G.T. 表示与真实图像计算FID指标。w/ Orig. 表示与1个设备原始实现生成的图像计算FID指标。
表2: Pixart和Flux.1上并行方法的FID分数。w/ G.T. 表示与真实图像计算FID指标。w/ Orig. 表示与1个设备原始实现生成的图像计算FID指标。
  • 视觉比较 (附录A, 图10):视觉对比显示,PipeFusion生成的图像与原始模型生成的图像在人眼看来几乎没有区别。

A7 补充细节

混合并行方案:尽管本文证明了PipeFusion相对于张量并行和序列并行的优势,但我们承认它并非通用解决方案。实际上,将PipeFusion与现有方法结合成混合并行方案可能是最实用的部署方向。这样的设计能够实现DiT推理在大型跨节点配置上的高效扩展:序列并行可以在NVLink连接的高带宽环境中使用,而PipeFusion在跨节点的低带宽环境中尤其有效。

对扩散步数的依赖:PipeFusion的一个关键限制在于它对多个扩散步骤的依赖。对于只有少数几个步骤的模型——例如像4步Flux.1-schnell【索引编号31,black-forest-labs/flux.1-schnell · hugging face,2024】这样的蒸馏变体或单步扩散方法【索引编号32,One-step diffusion with distribution matching distillation,2024,Proceedings of the IEEE/CVF conference on computer vision and pattern recognition;索引编号33,Mean flows for one-step generative modeling,2025,arXiv】——时间冗余极小,使得PipeFusion效果不佳。然而,在实践中这并不构成主要限制,因为最先进的DiT模型通常需要大量的步骤来提供有竞争力的质量。

在视频生成领域的应用前景:将PipeFusion应用于视频生成是另一个有前景的扩展方向。从架构上看,PipeFusion已经与多模态DiT(MM-DiT)如Flux.1和Stable Diffusion 3兼容,在这些模型中,跨模态条件通过统一的注意力机制注入。最近的视频模型(如Wan2.1【索引编号34,Wan: Open and advanced large-scale video generative models,2025,arXiv】、HunyuanVideo【索引编号35,Hunyuanvideo: A systematic framework for large video generative models,2024,arXiv】、CogVideo【索引编号36,Cogvideo: Large-scale pretraining for text-to-video generation via transformers,2022,arXiv】)采用类似的MM-DiT主干,这表明PipeFusion只需稍作调整即可应用。从张量形状的角度看,图像和视频DiT都操作于形状为(B, L, H)的潜在张量,视频模型将空间和时间维度展平为一个统一的序列维度。因为PipeFusion是沿着序列维度进行分区的,这种结构差异不影响其适用性。从性能角度看,视频模型通常更大,涉及更长的序列长度,这会放大SP的通信开销。PipeFusion通过参数分区和流水线执行来缓解这种开销,提供了更低的内存占用和更强的可扩展性。这表明视频生成不仅是应用PipeFusion的自然领域,也是一个非常有利的领域。

A5 结论

我们介绍了PipeFusion,一种专为扩散变换器设计的补丁级流水线并行策略。通过重用时间上冗余的激活值并重叠通信与计算,PipeFusion在保持图像质量的同时,降低了带宽需求和内存占用。在Pixart、Stable Diffusion 3和Flux.1上的实验证明,与强大的基线相比,PipeFusion在商用多GPU集群上实现了高达1.5倍的推理速度提升,取得了持续的延迟改进。

除了效率提升,PipeFusion还揭示了一个更广泛的系统层面的见解:扩散过程中的时间冗余可以被系统地利用来设计新形式的并行计算。虽然我们的研究集中在单节点、多GPU的推理上,但该方法可以很自然地与现有的序列或张量并行相结合,为在大型集群上进行可扩展部署提供了一条路径。我们希望这项工作能激励对混合并行策略的进一步探索,以弥合快速增长的生成模型与大规模服务它们的实际需求之间的差距。

A6 附录

A 图像质量

图像质量视觉评估:在图10中,我们选择了一些使用DistriFusion【索引编号16,Distrifusion: Distributed parallel inference for high-resolution diffusion models,2024,Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)】中的提示生成的样本图像,以帮助读者直观地评估生成质量。在各种补丁数量和设备数量的配置下,由PipeFusion生成的图像与原始图像在人眼看来几乎无法区分。请注意,FID评估结果与用于视觉演示的具体提示无关。

图10: 上图:使用Pixart的PipeFusion和DistriFusion生成1024px图像的展示。下图:对Pixart和Flux.1评估的FID分数。我们对DistriFusion和PipeFusion使用20步DPM-Solver,预热步骤设为1。我们使用COCO Captions 2014 [37]数据集评估FID分数。评估期间,从验证集中采样一个包含30,000张图像的子集,并调整大小至256px作为参考数据集。同时,每个实验生成30,000张256px的图像,每张图像与来自COCO Captions 2014数据集的标题配对,作为样本数据集。无论使用4个还是8个设备,以及在不同的补丁数量下,PipeFusion生成的图像质量都与原始图像非常接近。图像上方的FID是使用Clean-FID [38]与真实图像计算得出的。
提示:一辆双层巴士在街上行驶。
提示:一只棕色的狗躺在地上,面前有一个金属碗。
提示:一只多彩的鹦鹉把脚举到它的喙上。
提示:一个戴着耳机用笔记本电脑的孩子。
提示:一对反映已过期时间的停车计时器。

B 消融研究

分块数量 (Patch Number):我们分析了补丁数量M的设置对PipeFusion的影响,如图11所示。根据表1的分析,M不影响通信成本。然而,如图所示,当M等于GPU数量N时,延迟最低。如果M很大,算子的输入尺寸会变小,这可能对算子的计算效率产生负面影响。相反,如果M很小,重叠通信和计算的效率会降低。

图11: 在使用Pixart的三种不同分辨率的图像生成任务中,PipeFusion(无预热)在2x、4x和8xL40 GPU上针对不同补丁数量M的延迟。
图11: 在使用Pixart的三种不同分辨率的图像生成任务中,PipeFusion(无预热)在2x、4x和8xL40 GPU上针对不同补丁数量M的延迟。

预热步数 (Warmup Step):预热步骤的数量会对PipeFusion的整体性能产生负面影响。在最初的几个扩散时间步中,输入时间冗余相对较小。因此,通常采用几个预热步骤来启用同步通信,而不依赖于陈旧的激活值,这会给流水线带来等待时间,通常称为气泡。如表3所示,随着采样步骤数量的增加,延迟也会增加。如果预热步骤的比例相对较小,则无需采取任何措施。我们的实验涉及的采样步骤数量相对较少,具体为20和28步。在其他相关工作中,通常使用50步【索引编号16,Distrifusion: Distributed parallel inference for high-resolution diffusion models,2024,Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)】甚至100步【索引编号12,Deepcache: Accelerating diffusion models for free,2024,Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition】,这有助于减轻预热的开销。

表3: 预热步骤数对Pixart、sd3-medium和Flux.1-dev在1024px图像生成上的影响,其中pp表示pipefusion(pp)并行度。
表3: 预热步骤数对Pixart、sd3-medium和Flux.1-dev在1024px图像生成上的影响,其中pp表示pipefusion(pp)并行度。

缓解预热性能下降的策略:为了减轻预热引起的性能下降,我们可以将预热步骤与剩余的工作步骤分开,并为它们分配不同的计算资源。预热步骤后的输出特征图可以从预热设备传输到工作设备。预热阶段可以应用序列并行来充分利用计算资源。

预热步数的模型依赖性与未来工作:我们发现预热步骤的设置高度依赖于模型。通过将KV缓冲区初始化为零并且不执行预热,Pixart仍然可以实现令人满意的图像生成结果。通过动态检测潜在空间输入与前一个扩散步骤之间的差异,我们可以自动设置预热步骤。这些对预热的优化将是我们未来的工作。

在NVLink连接的GPU上的评估:为了评估PipeFusion对硬件互连带宽的敏感性,我们在通过NVLink(高带宽)连接的8×A100 GPU上进行了额外的实验,与默认的8×L40 PCIe设置形成对比。表4报告了PixArt在不同分辨率下的推理延迟。PipeFusion在所有分辨率下始终实现最低延迟,并且即使在高带宽的NVLink设置下也保持稳健。虽然由于计算成本的主导地位日益增强(即从通信密集型转向计算密集型),相对性能差距在更高分辨率下有所缩小,但PipeFusion仍然具有竞争力或优于所有基线。我们还强调,FID分数与GPU类型或互连带宽无关,因为它仅由模型架构和推理算法决定。

表4: 在8×A100 GPU上使用NVLink互连的PixArt延迟(秒)。PipeFusion始终提供最低的延迟。
表4: 在8×A100 GPU上使用NVLink互连的PixArt延迟(秒)。PipeFusion始终提供最低的延迟。