PipeFusion：如何用PCIe互联GPU 低成本并行推理扩散模型

今年二月Sora横空出世，我们正在跑步进入视频生成时代。同时，Sora巨大的部署难题也引爆了长序列的DiT推理方法研究热潮。面对这个问题，我的团队最近在这这方面做了一个非常有趣的工作叫PipeFusion，它可以显著降低DiT模型并行推理的带宽需求，能在PCIe互联的GPU上更有性价比地部署 DiT并行推理。希望它也可以为更大规模并行推理提供了一个新的思路。

论文地址如下，实验代码也开源了：

https://arxiv.org/abs/2405.14430

这个工作的缘起是几个月前听一次采访贾扬清老师Podcast，他提到最近和MIT和韩松，Medusa作者Tianle Cai一起中了一篇CVPR，里面设计了一个并行扩散模型推理的方法。我立刻去拜读了一下，读完之后大受启发，觉得扩散模型有一些独特的性质，可以在System方面加以利用。因此，我们在DistriFusion的一个关键发现基础上，提出了一个通信和内存方面性质更好并行方法PipeFusion。DistriFusion还是针对U-Net backband的扩散模型，顺应潮流，PipeFusion也将研究目标转移到了DiT。欢迎感兴趣的朋友，洽谈合作。

https://arxiv.org/abs/2402.19481

Sora时代Diffusion Model多卡推理需求

扩散模型（Diffusion Models）是图像和视频合成的首选技术。为了生成高分辨率图像和长时间的高保真视频，扩散模型也在演进。一方面，基础网络架构正在从传统的卷积神经网络U-Net转向Diffusion Transformers（DiTs），因为后者具有更高的模型容量和可扩展性。另一方面，输入序列长度在不断增长，比如大家预估Sora的上下文的visual tokens数量至少1M，才能达到高度的内容一致性。

DiT模型在生成长视觉序列时的推理延迟显著增加，例如，Sora模型的推理延迟可达数分钟，单个GPU无法满足实际应用的延迟要求。因此，DiT模型的部署必然是多卡并行处理。Attention机制导致的计算时间与序列长度呈二次增长，为了达到同样延迟指标，GPU数应该随着序列长度平方项增长的，所有部署Sore很大概率需要多机多卡大集群执行单个视频生成，DiT的并行推理扩展性至关重要。

尽管大型语言模型（LLMs）中常用的并行化技术，如张量并行和序列并行，通信量和激活参数正相关，在扩散模型中由于大激活小参数特点，这些技术效率不高。通信成本往往超过了并行计算的好处。因此，DiT模型的部署仍然需要配备高带宽互连（如NVLink+RDMA）的GPU集群。这样的话，部署DiT的成本就居高不下了。

Diffusion Model推理原理和特性

尽管DiT和LLM的结构相似都是Transformers架构，看起来工程挑战类似。这里有必要给大家介绍一下扩散模型推理原理。

扩散模型的训练过程：给一个图片，经过很多步骤，每一步加噪声最后变成一个全是噪音的图片。训练过程就是预测一个Noise Predictor的监督学习任务，输入是中间状态图片、输入指令和timestep来ground truth就是该timestep的噪声。

扩散模型的推理过程：给一个噪声，通过Noise Predictor来通过多次去噪，最后变成一个有意义的图片。每一次去噪就是一个Diffusion Step。

通过描述可知，扩散模型推理和LLM的有一个很大的差异点，就是扩散模型是重复计算很多相同的Diffusion Step，而且连续的Diffusion Step之间输入数据和激活状态之间存在的高度相似性。我们称之为输入时间冗余（Input Temporal Redundancy）。这种冗余性表明，在进行相邻时间Diffusion Step去噪时，可以利用前一步骤的过期Activation来顶替当前步骤生成的Activation做一些计算，可能结果也不会差太多。

最先利用这个特性来加速扩散模型并行推理的是MIT韩松和Lepton贾扬清合著的论文DistriFusion。他们指出这种特性被用来优化并行推理过程，通过分割输入图像为多个区域，并分配给不同的计算设备处理，同时利用局部fresh activation和one step stale activation混在一起进行Attention和Convolution操作最终生成效果近似，从而可以隐藏通信成本，实现很好的并行。但是我们发现，可能会导致内存使用效率低下，因为需要维护完整的注意力键（K）和值（V）的空间形状，这限制了其在多计算设备环境下的可扩展性。

Diffusion Model并行方法

在DiT推理中，鉴于结构的相似性，我们可以用LLM中非常成熟的张量并行（TP）和序列并行（SP）。也可以使用一些利用扩散模型Input Temporal Redundancy特性的并行方法，比如DistriFusion和我们提出的PipeFusion。

下表对比了不同DiT并行方法的内存和通信性能。表中，p 表示生成序列的长度（latent space中的像素数），hs 代表模型的hidden size。L 是网络层数，P 是总参数数量。注意力模块中Query（Q）、Key（K）、Value（V）和Output（O）的参数数量相同，表中用 A 表示，峰值Activation内存与单个DiT层的成本相同，就是attention的峰值Activation，即Q、K、V、O的总量。在attn-KV列中，fresh表示使用当前扩散步骤的KV Activation，而stale表示使用前一步骤的部分KV Activation，stale-表示更偏向fresh。名称后的 * 表示该方法采用异步通信，通信成本可以通过计算隐藏。

1. 张量并行（TP）:

张量并行通信量是最大的，而且是同步的，一次正向传播每次需要等价2此AllReduce通信。它的好处参数和激活的需求都减少到1/N，对内存很友好。

2. 序列并行（SP）:

对扩散模型latent space的图片切分也就是LLM中对序列维度切分。序列并行通信量比较复杂，Ulysses和TP一样，因为4次All2All等价于2次AllReduce；Ring方式是可以重叠的，但会减慢Attention计算。SP的参数内存需求不变，但是激活的需求都减少到1/N，所以内存是不如TP的。

3. DistriFusion：Displaced Patch Parallelism

DistriFusion是一种序列并行，它也将输入图像分割成多个Patch，并分配给不同的计算设备处理。它和序列并行区别在于，KV Activation的通信模式不同。DistriFusion使用前一个diffusion timestep的部分的stale K, V与当前步骤的部分的fresh K，V相结合，来执行注意力操作。Stale KV通信可以和当前diffusion timestep forward计算隐藏，但缺点是需要维护一个较大的通信Buffer，而且其内存成本不会随着计算设备的增加而减少。

我们将DistriFusion和同样重叠KV通信和计算的Ring-Attention进行对比，二者都可以重叠通信。如下图所示，DistriFusion需要跨越Diffusion Step维护完整的KV，而Ring把KV进行了切分。所以，DistriFusion可以用更大范围的计算（一次完整网络forward）重叠KV通信，而Ring-Attention必须在当前Attention算子的计算重叠KV通信。所以，DistriFusion可以看成用更多内存为代价的异步通信方式的序列并行。

4. PipeFusion：Displaced Patch Pipeline Parallelism

有了上述工作的铺垫，大家就更容易理解我们工作的价值。PipeFusion是一种更好利用扩散模型Input Temporal Redundancy的特性的方法。和PipeFusion不同，它通过流水线方式组织异步通信和计算。

对比DistriFusion和PipeFusion

PipeFusion将输入图像分割成_M_个不重叠的Patch，并将DiT网络均匀分成_N_个阶段，每个阶段由不同的计算设备顺序处理。切分与序列并行和DistriFusion略有不同，这里的M和N可以不相等。每个设备以流水线方式处理其分配阶段的一个Patch。这种需要均分网络的方法非常适合DiT模型，因为它们包含许多重复的Transformer块。比如，U-Net扩散模型没有这种重复结构，切分阶段就很费劲。

如下图所示，以_N_ = 4和_M_ = 4为例，展示了两个diffusion timestep的流水线工作流程，其中步骤T+1（浅灰色）的计算在步骤T（深灰色）之前完成。例如，Patch 0（P0）在deviuce 0上计算后传递给device 1，同时device 0并行计算P1（Patch 1）。由于Sampler的计算是逐元素操作，所有diffusion timestep都可以流水线化。利用输入时间冗余性质，设备无需等待接收当前流水线步骤的全空间形状KV Activation即可开始其阶段的计算，它使用前一步骤的stale activations代替fresh activations为当前步骤所用。因此，一旦流水线初始化，流水线内就没有等待时间。考虑到流水线中的气泡，流水线的有效计算比率也非常高，因为diffusion timestep通常是几十甚至几百。

注意，PipeFusion和LLM中流水线并行还是有很大区别的。LLM流水线并行训练或推理要求输入是一个batch的序列，而我们不需要输入是一个batch的图片，而是在一个图片生成任务上做流水。

PipeFusion在通信上仅传输不同stage之间的输入和输出激活（通常命名为hidden states），而不是像DistriFusion一样传递每层的K，V。hidden states和K或V基本大小一样。如上表所示PipeFusion通信量没有_L_项，相比其他方法这就极大减少了通信开销。并且，PipeFusion通信和计算可以重叠，发送前一个Patch计算和接受后一个Patch计算都可以和当前Patch计算使用异步P2P重叠。而且，PipeFusion对网络拓扑要求很低，它的通信呈一个环状，不需要所有计算节点全互联。

在内存使用上，每个设备仅存储整体网络1/N的参数和1/M的QO激活，KV激活需要维护L/N个层的全量KV，extra buff随着设备数增多而成比例减少，这比DistriFusion内存效率显著提高。

不仅内存和通讯特性上更好，理论上PipeFusion在生成图像的准确性上理论上优于DistriFusion，因为它使用了更多的fresh activation信息。比如，当M=4时，一个Diffusion Timestep内不同时刻的fresh vs stale Activation的比例如下，深灰色是fresh activation，浅灰色是前一步的stale activation，可见fresh部分逐渐增大。而DistriFusion在整个Diffusion Timestep内一直只有1/N部分是fresh的。

实验效果：

我们使用pixart-alpha模型，它是一种DiT模型变种，可以改变输出图片分辨率。我们在两种PCIe设备上，8xL20和4xA100上生成1024px到8192px图像。PipeFusion基本都获得了最低的延迟表现，同时能够胜任PipeFusion OOM的场景。

8xL20(PCIe) Latency评估

Figure 8展示了8xL20 (48GB) PCIe集群上的延迟。在1024px和4096px的分辨率下分别达到了其它最佳方法1.47倍和1.31倍的加速比。在1024px和4096px的场景下,相较于单GPU的Baseline,PipeFusion分别实现了2.46倍和4.3倍的加速比。注意，DistriFusion和SP在4092px分辨率生成任务OOM了。

4xA100(PCIe)Latency评估

在4卡A100 PCIe机器上，延迟指标展示在Figure 6。在1024px分辨率上比SOTA方法提高了2.01倍，在2048px分辨率上，比SOTA方法提高了1.48倍，在8192px分辨率上，比SOTA方法提高了1.10倍。注意，DistriFusion和SP在4092px分辨率生成任务OOM了。

内存使用展示在Figure 7，可见TP是内存最少的，但是它通信太多了。DistriFusion显著优于除了TP以外的方法。