TACO-DiT:面向图片/视频生成的DiT扩散模型实时推理解决方案
图片与视频生是大模型多模态应用的关键领域,它们使用DiTs为架构的扩散模型来处理超长输入序列。然而,由于Transformer的计算复杂度随序列长度呈平方级增长,导致DiTs在实际部署中的延迟问题尤为突出。以V100为例,利用HunyuanDiT生成一幅4096px的图片竟耗时长达1小时13分钟。TACO-DiT作为一款专为DiTs企业级部署设计的高性能推理框架,它可以将DiT推理并行扩展至多卡乃至多机超大规模设备,结合编译优化技术,确保应用的实时性需求得以满足。以最新发布的Flux.1模型为例,在A100设备上,TACO-DiT实现了一秒出图!即便对于混元DiT,也仅需五秒即可完成出图。此外,针对T4/V100等老旧GPU设备,TACO-DiT展现出卓越的高性价比优势,未来也计划支持公司自研紫霄V2芯片。TACO-DiT是我司完全自研的软件,其中PipeFusion,USP等并行技术已经在业界产生了广泛的影响力,引领风气之先。
本文篇幅较长,
希望一文中全面展示DiTs的进展和部署挑战,
读者可以各取所需,
如有合作需要请联系
>>欢迎扫描二维码,加入社区<<
1
2024年百花齐放视觉AIGC产业
2024年,多模态是大模型领域的主旋律,图片和视频生成为代表的视觉AIGC技术迎来了长足的发展。
今年二月份,Sora横空出世,让大家看到了视频生成的广阔前景。国外视频生成公司Luna、Pika、Runway等如火如荼,国内也不甘示弱。四月生数科技的ViDu项目成功创作出16秒的高质量电影级视频。六月,快手推出的可灵平台并对外开放使用,以其独特的老铁风格和高度逼真的生成效果,迅速获得了公众的认可,使得之前未公开测试的Sora平台被戏称为“美国可灵”。此外,阿里EMO、智谱清影、Minimax等公司的视频生成产品也在近期也相继发布。
在传统图像生成领域,新文生图模型不断涌现。混元DiT,Stable Diffusion 3相继问世,生成的图像质量更高,细节更加精细。其中我司混元DiT也是第一个开源的DiT模型,引领风气之先。刚刚过去的七月,Flux.1是由Black Forest Labs出。13B参数的DiTs Backbone接近Llama2-13B,也就说参数量级上赶上了一年前最常用的LLM模型。这也是开源扩散模型首次达到10B量级。Flux.1在性能上全面超越了Midjourney、DALL-E 3等竞品,支持多种分辨率和宽高比,能够生成细节丰富、质量上乘的图像。
这些进展不仅推动了AI图像生成技术的边界,也为未来的创意应用和商业解决方案提供了强大的可能性。随着这些技术的不断发展和优化,我们可以期待在不久的将来,AI图像生成将在艺术创作、影视制作、社交媒体,广告创意,室内设计、娱乐等领域提供了内容创作的有力工具。
下图是用Flux.1生成图片后,用视频生成软件做成视频。一段新闻播报片段栩栩如生。
2
扩散模型和Diffusion Transformers(DiTs)
应用的百花齐放来源于技术的进步,采用DiTs(Diffusion Transformers)结构的扩散模型是图像和视频生成的背后关键技术。我们本小节先对扩散模型和DiTs做分别介绍以便大家理解后续内容:
2.1 扩散模型(Diffusion Models)
扩散模型是一种生成模型的范式,是与LLM所采用自回归生成范式不同的另一种范式。扩散模型通过逐步向数据添加噪声,然后学习如何逆转这个过程来生成新的数据样本。正如开朗基罗的名言: 雕像本来就在石头里,我只是把不要的部分去掉。
扩散模型的训练过程分两个阶段:
前向加噪(Forward Diffusion):然后,通过逐步增加噪声的方式,将数据转换成一系列越来越模糊的样本。这个过程通常被称为“扩散过程”。如下图所示,小狗的图片最后变成了噪声。
反向去噪(Reverse Diffusion):接下来,模型需要学习如何逆转扩散过程,即从噪声数据中恢复出原始数据。这是通过训练一个神经网络Noise Predictor来实现,该网络学习如何预测每一步中应该减少多少噪声。如下图所示,神经网络的输入是红框部分,包括一张加噪声的图片,输入提示词或者图片等Conditioning,扩散迭代的步数。神经网络的Ground Truth就是当前步骤加入的噪声信号。
扩散模型的推理过程就是从高斯分布开始,用训练好的Noise Predictor生成噪声,逐步去噪生成有意义的图片。
直接对图片进行操作计算量太大,目前主流扩散模型都采用LDM方式:扩散模型先把图片从像素空间用VAE压缩到Latent Space,然后在Latent Space进行加噪和去噪。推理完成后,使用VAE从Latent Space解码到像素空间。
2.2 Diffusion Transformers(DiTs)
DiTs是扩散模型主干网络(上一节Noise Predictor)的一种网络结构。在2024年之前,扩散模型主干网络主要是基于 U-Net 架构发展的。
2023年,Meta提出利用Transformer处理扩散模型DiTs,通过将图像转化为token序列并结合标签嵌入,验证了类别控制图像生成的有效性。同年,华为的PixArt-alpha,清华的U-ViT也都验证了DiTs的有效性。随着2024年 OpenAI 发布 Sora 以及 Stability AI 发布 Stable Diffusion 3,基于 Transformer 架构的扩散模型 (如 DiT) 已经成为绝对主流。最新发布的Flux.1和腾讯的混元DiT都是基于DiTs而非U-Net架构。
如下图所示,以State Diffusion为代表的U-Net网络主要是使用卷积操作,而DiTs则采用和LLM类似Transformer结构。因为后者随着增加参数量,可以观测到模型效果的提升,也就是Scaling Law,因此正在逐渐取代U-Net成为扩散模型主干网络的事实标准。
3
长序列DiTs实时部署工程挑战
DiTs技术正蓬勃发展,目前已涌现出诸多优秀的DiTs模型。在文生图领域,混元DiTs、Flux.1等模型表现卓越;而在文生视频领域,CogVideoX、OpenSora等也展现出强大的实力。训练这些模型和LLM想来说已经是小菜一碟,但是如何有效部署这些模型,让它们真正在业务落地,仍存在巨大挑战。
其主要源于长序列所引发的巨大计算需求以及模型结构的多变性带来的软件工程需求。
第一,随着生成质量的提高,DiTs所需的输入序列长度也在增加,而Attention的计算量呈平方级增长,这将导致非常高的推理延迟。例如,在单卡V100上,1.5B的HunyuanDiT模型生成1024px的图片需要50秒,如果想生成4096px的图片则需要1小时13分钟(实测4402.25秒)。在高清图片生成上尚且如此,面对未来1M tokens序列长度起步的视频生成,即使是最先进的B200也会力不从心。Sora迟迟没有对外开发,其主要原因也是高居不下的推理延迟无法满足并发使用需求。据投资机构Factorial Funds估算Sora推理需求:一个H100 GPU一小时生成大约5分钟的视频。笔者曾经现场体验了Minimax的视频生成接口,生成720p 5秒钟视频要四分钟,让用户意兴阑珊。
为了满足实时的用户体验,多GPU乃至多机部署DiT已成为必然趋势。然而,目前市场上缺乏一个能够满足日益增长的DiT模型和实时性要求的通用推理框架。为此,我们从头开发了一个通用DiTs推理引擎TACO-DiT,旨在提供一个可扩展、高性能且易于使用的解决方案,以满足未来图片、视频生成应用的需求。
第二,和LLM收敛到以LLAMA为代表相对固定的模型结构不同,DiTs模型设计更像是传统CV模型,模型结构百花齐发,变化多端,每种模型结构给人感觉都是灵感与艺术的结合。如下图所示,Flux.1,HunyuanDiT,Pixart三个DiT模型结构存在肉眼可见的差别。
具体来说,如下模型设计的细节都存在很多种选择,这给模型的工程优化带来了巨大的困难,常常出现刚刚针对HunyuanDiT进行优化,发现切换成SD3却又要重头来过。
DiT Blocks之间连接方式:Skip-Connection vs Linearized
Conditioning信息注入方式:In-contex Conditioning, Cross-Attention, adaLN-Zero
位置编码: RoPE,固定位置编码
Attention实现:Parallel Attention, Approximate Attention
效果增强技巧:Classifier Free Guide
4
TACO-DiT概览:DiT推理的TensorRT-LLM
与LLM已具备成熟的推理解决方案,如vLLM和TensorRT-LLM不同,DiTs结构目前尚未形成一个通用的框架,以满足并行推理和多模型适配的广泛需求。过去几年,扩散模型的传统部署主要面向个人主机,例如装配3080显卡,使用方式主要是通过webui或comfyui图形界面进行单卡推理,而未曾考虑企业级的需求。然而,随着DiTs模型能力的提升,企业化需求开始初现端倪。例如,12B的Flux.1模型发布后,Block Forest Lab并未开源其能力最强的Flux.1 Pro版本,而是将其授权给Fal.ai和Replicate.ai进行私有化部署。此外,以可灵为代表的视频生成应用也非个人发烧友可以驾驭,视频生成模型后续在企业级生产场景的工程化部署,需要多机多卡并行运算的推理环境。
瞄准即将爆发的市场需求,腾讯云开发了TACO-DiT,这是一个满足企业级实时推理部署DiTs模型的高性能推理引擎。类比TensorRT-LLM之于LLM,我们希望TACO-DiT成为DiT领域的TensorRT-LLM,其核心能力如下所示:
第一,DiT主干网络和VAE并行能力:TACO-DiT可以使用多种并行混合方式扩展到非常大的规模,比如多机多卡来实现良好的扩展性。其中包括腾讯原创技术混合序列并行USP和异步流水并行PipeFusion。
第二,编译优化技术:面对多变的网络结构,我们使用torch.compile技术对GPU核心进行加速。
第三,多元异构算力兼容:TACO-DiT支持多种NVIDIA GPU,很多古早卡型我们也精心适配,包括H800/A800/L40/L20/ V100/T4等。我们计划支持我们的自研芯片紫霄V2,目前已有一些初步成果。
第四,易用的使用方式:TACO-DiT可以支持原生diffusers模型,也支持以comfyui方式对模型进行部署。我们同时也在腾讯云HAI上提供服务。
5
TACO-DiT的核心技术
TACO-DiT是腾讯完全自主研发推理引擎,因为没有业界同类方案参考,每一行代码都是我们从头搭建的。TACO-DiT以创新的技术解决了如下难题。
5.1 主干网络混合并行
并行推理是TACO-DiT最核心能力。如图所示,TACO-DiT提供了一套全面的并行化技术。
LLM常用的张量并行并不适用于DiT推理,因为DiT序列长度尤其的长,通信量也更大,对通信能力要求更高。
对于DiT主干,它提供四种基础方法,数据并行、混合序列并行USP、PipeFusion和CFG并行(途中黄色部分),以混合方式运行。此外,张量并行和DistriFusion并行方法也可各自独立发挥作用。对于VAE模块,TACO-DiT提供了防止内存溢出(OOM)问题的DistVAE并行实现。图中(TACO)强调了腾讯首次提出的方法。
对于主干网络,TACO-DiT支持四种基础并行策略方法:
PipeFusion:腾讯云首创的异步流水线并行方法,在弱互联,比如pcie/以太网,的网络硬件上有优势,比如以太网互联的多机多卡情境下效率远高于其他方法。
Sequence Parallel:混合序列并行,使用腾讯云首创的ring & ulysses sequence parallel混合序列USP。
Data Parallel:在输入多个prompt或单个prompt生成多张图片时,在image间并行处理。
CFG Parallel, a.k.a Split Batch:在模型使用classifier free guidance(CFG)时可以开启,并行度恒定为2。
TACO-DiT中这四种并行方式可以以任何形式排列组合,混合策略在通信量和通信拓扑两方面可以取得最佳效果,使TACO-DiT可达到非常高扩展性。
5.2 PipeFusion:异步序列流水并行
PipeFusion是一种TeraPipe式的序列流水并行。但和针对LLM设计TeraPipe有问题有显著不同,它利用了Input Temporal Redundancy——即扩散步骤间输入与激活之间的数值相似性,这是扩散模型特有的特性,同样被应用于DistriFusion中。Input Temporal Redundancy和序列流水并行结合,避免了LLM模型中因果注意力带来的负载均衡问题。PipeFusion可以显著减少了通信量,它和序列并行、张量并行、DistriFusion对比都是通信带宽需求都是最低的,从而特别适合在以太网连接的多节点设置和PCIe连接的多GPU上使用。
下面我具体介绍PipeFusion算法的细节,感兴趣同学可以阅读我们公开的论文。
PipeFusion: Displaced Patch Pipeline Parallelism for Inference of Diffusion Transformer Models
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通信量和网络层数无关,相比其他方法这就极大减少了通信开销。并且,PipeFusion通信和计算可以重叠,发送前一个Patch计算和接受后一个Patch计算都可以和当前Patch计算使用异步P2P重叠。而且,PipeFusion对网络拓扑要求很低,它的通信呈一个环状,不需要所有计算节点NVLink全互联。因此也深受国产硬件厂商所喜爱,比如燧原、天数都选择PipeFusion进行多卡并行。
因为利用了Input Temporal Redundancy特性,有人可能担心PipeFusion对生成精度的影响。论文中我们也有详细实验评估PipeFusion对精度影响。我们使用FID指标(越低越好)来评估精度,PipeFusion和原始串行算法表现类似,而且整体上略低于DistriFusion,这也符合我们的预期。同时,肉眼也难以分辨生成结果和单机版本的差异。
5.3 2D序列并行USP
因为模型小而输入长的特点,序列并行非常适合DiT推理。DeepSpeed-Ulysses和Ring-Attention是目前两种主流的序列并行方法。但这二者并不是非此即彼的二选一,我们团队首先提出二者应以混合方式共同切分序列维度,可以调节二者并行度保证ulysses-degree X ring-degree=sp-degree即可,我们称它为Unfied-SP(统一序列并行,USP),或者业界通常称之为2D序列并行或者混合序列序列并行。USP这带来的好处是,它覆盖原来二者的能力,没有任何损失,只有额外好处。首先,避免了Ulysses sp-degree必须小于attention head的限制,且混合并行的通信模式对异构的网络更有好,在PCIe和多机多卡环境上相比单独使用Ulysses和Ring都有加速。
USP不仅作为TACO-DiT序列并行方法,还可以用在LLM的训练和推理中,短时间内产生了非常广泛的影响。很多友商也在使用USP,我们发现阿里云(宣传稿)、华为昇腾(宣传稿,MindSpeed开源)、上海AI实验室(论文主动合作)、智源实验室(FlagScale开源框架)、咱们腾讯的混元团队(论文引用)、NVIDIA美国(论文+开源软件VILA)公开宣布使用。
感兴趣同学也请参考我们的技术报告。
USP: A Unified Sequence Parallelism Approach for Long Context Generative AI
5.4 PipeFusion+序列并行混合并行
TACO-DiT设计目标是扩展DiT推理过程到超大规模的计算设备上,比如异构网络互联条件下多机多卡,比如以太网和PCIe。单一并行方式,比如PipeFusion或者序列并行,很难同时做到这两点,不同并行方式混合在一起变得尤为重要。
TACO-DiT支持四种并行方式:PipeFusion、Sequence、Data和CFG Parallel。其中,Data和CFG Parallel都是在图像间并行相对简单,而PipeFusion和Sequence在图像内部的不同Patch间并行则较为复杂。能让这两种并行方式的混合使用,正是TACO-DiT核心创新点之一。
PipeFusion利用Input Tempor Redundancy特点,使用过时的KV(Stale KV)进行Attention计算。也就是说,参与Attention的KV并不是和输入切分方式一样。比如输入是1/N的序列,如果只使用1/N KV,那就是经典的混合并行,和Megatron的DP+TP使用方式一样。但是,TACO-DiT中参与计算的KV确是完整的KV,这导致使用标准的序列并行接口,如RingAttention、Ulysses或USP是不能胜任PipeFusion+USP混合并行。
正确通信有效的KV数据并不简单,我们对这个问题具体说明。下图展示了pipe_degree=4,sp_degree=2的混合并行方法。进程组划分如下。
设置num_pipeline_patch=4,图片切分为M=num_pipeline_patch*sp_degree=8个Patch,分别是P0~P7。标准序列并行Attention接口实现,输入Q,K,V和输出O都是沿着序列维度切分,且切分方式一致。如果不同rank的输入patch没有重叠,每个micro step计算出fresh KV更新的位置在不同rank间也没有重叠。如下图所示,standard SP的KV Buffer中黄色部分是SP0 rank=0拥有的fresh KV,绿色部分是SP1 rank=1拥有的fresh KV,二者并不相同。在这个diffusion step内,device=0无法拿到P1,3,5,7的fresh KV进行计算,但是PipeFusion则需要在下一个diffusion step中,拥有上一个diffusion step全部的KV。standard SP只拥有1/sp_degree的fresh kv buffer,因此无法获得混合并行推理正确的结果。
TACO-DiT为PipeFusion+序列并行混合设计了专属的序列并行Attention接口,以适应这种混合并行的需求。TACO-DiT使用图中Hybrid-SP把SP的中间结果存在KV Buffer内。这样效果如下图,每个micro-step SP执行完毕后,SP Group内不同rank设备的fresh KV是replicate的。这样一个diffusion step后,SP Group所有设备的KV Buffer都更新成最新,供下一个Diffusion Step使用。
5.5 选择最佳混合并行策略
有了上述四种基础并行进行混合并行,如何选择最佳的混合并行策略?实践中,我们可以遍历不同的并行策略组合找到最佳的方式。不过,这里咱们理论分析比较一下不同的并行方法,从中读者也可以明白为什么我们选择PipeFusion+序列并行作为图片内并行策略,而没有选择张量并行和PipeFusion。
下表详细列出了DiTs中上述图像内并行性的通信和内存成本,除了CFG和DP(它们是图像间并行)。(*表示通信可以与计算重叠。)
𝒑: 像素数量,序列长度;𝒉𝒔: 模型隐藏层大小;𝑳: 模型层数;𝑷: 模型总参数;𝑵: 并行设备(GPU)数量;
𝑴: 补丁分割数量;𝑸𝑶: 查询和输出参数计数;𝑲𝑽: KV激活参数计数;𝑨 = 𝑸 = 𝑶 = 𝑲 = 𝑽: Attention的Q、K、V、O张量的参数相等;
我们可以得出一下结论:
1. PipeFusion拥有最低的通信成本,它的communication cost一项和层数L无关,而且P2P方式通信可以和计算隐藏,使之成为通信开销最小的并行方法。
2. 序列并行(Ring,Ulysses)的通信带宽需求小于张量并行,Ulysses的带宽需求随着GPU数增多而减少。另外,Ring的通信量也是张量并行的一半,而且Ring和DistriFusion相当,二者都可以重叠通信和计算。因为张量并行通信带宽显著逊于其他方法,因此我们没有将它纳入基础并行方法中。
3. 之所以没有在基础并行中添加DistriFusion,因为它的内存消耗极大,extra buff memory为2/N*A,每个机器都需要存储全部的KV,因此如果序列长度增大,内存很容易OOM。
5.5 VAE并行
由diffusers采用的stabilityai/sd-vae-ft-mse在处理高分辨率图像(A100上的8192像素)时会导致内存溢出(OOM),这一关键问题在diffusers/issues/5924中有详细记录。
为了解决这一限制,我们开发了并行VAE,这是一种能够并行高效处理高分辨率图像的解决方案。我们的方法结合了两种关键策略:
Patch Parallel:TACO-DiT将latent space中的特征图分割成多个块,并在不同设备上进行序列并行的VAE解码。这种技术将中间激活所需的峰值内存减少到1/N,其中N是所使用的设备数量。对于VAE中的卷积操作,我们需要传输图像的halo区域数据。
分块输入处理:类似于MIT-patch-conv,TACO-DiT将输入特征图分割成块,并依次将它们送入卷积操作,它最小化了Activation内存消耗。最近智谱发布的CogVideoX中3D也采用了类似的技巧。
通过协同这两种方法,TACO-DiT显著扩展了VAE解码的能力。我们的实现成功处理了高达10240像素的图像分辨率——与默认VAE实现相比,11倍增长。
现在很多文生图模型已经支持超高分辨率,比如Pixart-Sigma生成可以达到4096px,我们利用TACO-DiT探索了一下Flux.1 schnell版本的超分生成能力。原来超过2048px就会OOM,使用通过Parallel VAE让我们得以一窥更高分辨率生成能力的真容。
prompt是"A hyperrealistic portrait of a weathered sailor in his 60s, with deep-set blue eyes, a salt-and-pepper beard, and sun-weathered skin. He’s wearing a faded blue captain’s hat and a thick wool sweater. The background shows a misty harbor at dawn, with fishing boats barely visible in the distance."
生成2048px,3072px和4096px生成质量如下。2048px的生成质量还是很高的,但是3096px大部分图片都是虚化的,而4096px生成已经质量很低了。
5.6 编译加速
TACO-DiT采用两种编译加速技术,torch.compile和onediff,以提高GPU上的运行时速度。这些编译加速与并行化方法相结合使用。
由于torch.compile是PyTorch 2.0新功能,简单使用torch.compile并不能得到正确的结果。TACO-DiT对V100/A800/L40不同卡上进行了适配。将影响计算正确性的部分剔除出compiler优化范围。
我们暂时没有采用TensorRT,混元DiT的TensorRT也可以和TACO-DiT正交使用。
6
TACO-DiT案例效果
TACO-DiT已经支持如下图片和视频生成模型,我们挑选目前最流行的文生图模型Flux和腾讯混元展示加速效果。实验的baseline就是diffusers仓库的单卡标准推理实现。
6.1 HunyuanDiT
我们测试使用开源的HunyuanDiT,采样步数50步。
在8xA100(NVLink)机器上,在使用不同GPU数目时,最佳的并行方案都是不同的。这说明了多种并行和混合并行的重要性。最佳的并行策略在不同GPU规模时分别是:在2个GPU上,使用ulysses_degree=2;在4个GPU上,使用cfg_parallel=2, ulysses_degree=2;在8个GPU上,使用cfg_parallel=2, pipefusion_parallel=4。
torch.compile带来的加速效果也很可观,同样并行方案有1.26x到1.76x加速效果,对于8 GPU的场景是最明显的,有1.76x加速。
在8xL40 (PCIe)上的延迟情况如下图所示。同样,不同GPU规模,最佳并行策略都是不同的。和A100上不同,在L40上,8 GPU和4 GPU的延迟没有明显变化。我们认为是PCIe导致跨socket之间通信带宽过低导致的。
torch.compile带来1.2x到1.43x加速。
在8xV100上的加速下如下图所示。我们在8卡上将延迟降低到了10秒,torch.compile带来1.10x到1.30x加速。
在4xT4上的加速下如下图所示。
6.2 Flux.1
Flux.1开源进入diffusers代码仓库后,TACO-DiT也第一时间进行了支持。目前支持USP混合序列并行+VAE Parallel部分能力。PipeFusion还在开发中,CFG Flux.1没用,所以没法使用。我们对Flux.1 schnell版本进行测试,采样步数4步。
在8xA100 (80GB) NVLink互联的机器上,USP最佳策略是把所有并行度都给Ulysses,生成1024px图片仅需0.82秒,做到了一秒出图!生成2048px图片仅需2.4秒,相比单卡加速4.6x。我们对更高分辨率图片生成加速效果更明显,比如4096px达到7.4x,但值得注意,上面章节整理了Flux.1 4096px是没有意义的。
torch.compile起到了良好的加速效果。
在PCIe Gen4互联的8xL40上,4卡规模TACO-DiT也有很好的加速。生成1024px图片,使用ulysses_degree=2, ring_degree=2延迟低于单独使用ulysses或者ring,1.41秒可以生图。8卡相比反而会变慢,这是因为这是需要经过QPI通信。我们预期使用PipeFusion会提升8卡的扩展性。
我们在1024px图片生成任务上,对比torch.compile和onediff的性能差别。在1,8 GPU上,torch.compile略好,在2,4 GPU上onediff略好。
8xL40上生成2048px图片性能如下图。因为计算和通信比例增大,所以和1024px任务不同,8卡相比4卡延迟更低,生图最快可达3.67秒。
6.3 SD3
我们是用开源版本的stable-diffusion-3-medium-diffusers 2B模型进行性能评测。
在8xA100(NVLink)机器上,在使用不同GPU数目时,最佳的并行方案都是不同的。这说明了多种并行和混合并行的重要性。最佳的并行策略在不同GPU规模时分别是:在2个GPU上,使用cfg_parallel=2;在4个GPU上,使用cfg_parallel=2, pipefusion_parallel=2;在8个GPU上,使用cfg_parallel=2, pipefusion_parallel=4。
torch.compile在除了8 GPU的场景下都来来了加速效果。
在8xL40 (PCIe)上的延迟情况如下图所示。同样,不同GPU规模,最佳并行策略都是不同的。torch.compile都来了加速效果。
6.4 生数科技实际案例
今年五月PipeFusion一经提出,就被文生视频公司生数科技在线上部署所使用。具生数反馈,使用PipeFsuion相比之前使用Ulysses的序列并行方案,36%-52%(PCI-E 4.0, 4090),16%-18% (PCI-E 5.0, 4090)。
>>欢迎扫描二维码,加入社区<<