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。

image.png

通过描述可知,扩散模型推理和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。名称后的 * 表示该方法采用异步通信,通信成本可以通过计算隐藏。

image.png

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可以看成用更多内存为代价的异步通信方式的序列并行。

image.png

4. PipeFusion:Displaced Patch Pipeline Parallelism

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

image.png
对比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的图片,而是在一个图片生成任务上做流水。

image.png

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的。

image.png

实验效果:

我们使用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了。

image.png

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以外的方法。

image.png

精度评估

我们也有FID指标(越低越好)来评估精度,它和DistriFusion变现类似,整体上还是略低于DistriFusion,这也符合我们的预期。同时,肉眼也难以分辨生成结果和单机版本的差异。

image.png

总结:

本文介绍了PipeFusion,用于在多个设备上并行化DiT的推理。通过利用跨diffusion step的输入相似性,并采用流水线方法来协调通信和计算,PipeFusion减少了通信带宽和内存需求,使得DiT推理在PCIe连接的设备上更加高效。

The End

作者:方佳瑞
来源:GiantPandaCV

推荐阅读

欢迎大家点赞留言,更多Arm技术文章动态请关注极术社区嵌入式客栈专栏欢迎添加极术小姐姐微信(id:aijishu20)加入技术交流群,请备注研究方向。

推荐阅读
关注数
18676
内容数
1260
嵌入式端AI,包括AI算法在推理框架Tengine,MNN,NCNN,PaddlePaddle及相关芯片上的实现。欢迎加入微信交流群,微信号:aijishu20(备注:嵌入式)
目录
极术微信服务号
关注极术微信号
实时接收点赞提醒和评论通知
安谋科技学堂公众号
关注安谋科技学堂
实时获取安谋科技及 Arm 教学资源
安谋科技招聘公众号
关注安谋科技招聘
实时获取安谋科技中国职位信息