腾讯发布Project xDiT！打造DiT并行推理的vLLM

2024-08-12 13:08

©PaperWeekly 原创 · 作者 | 方佳瑞

单位 | 腾讯

研究方向 | 机器学习系统

上半年，我们见证了国内视频生成领域的迅猛发展。四月份，生数科技的 ViDu [1] 成功生成了 16 秒的电影级视频，不仅在视频号上引发了热烈反响，更揭开了中国 Sora 们崭露头角的序幕。而到了六月份，快手可灵 [2] 的发布，以其逼真的老铁风格，让人一时间真假难辨，更难能可贵的是它面向公众开放使用，使得迟迟未敢公测的 Sora，被戏称为“美国可灵”。

与此同时，阿里 EMO [3]、智谱清影 [4] 等视频类产品也相继问世，共同推动了这一领域的繁荣发展。不能忽视的是，在传统的图片生成领域，SD3、Flux、Pixart 等等也持续发力，图片更加高清、细节更加逼真。

DiTs（Diffusion Transformers）作为文生图与文生视频的核心网络结构，是引领高清图像与长时一致性视频生成潮流之关键。

随着生成质量要求的不断提升，DiTs 的输入序列长度日益增长，序列增长导致Attention 机制的计算量也随之呈平方级膨胀，对推理延迟 SLA 的影响是灾难级的。在单卡 A800 上 0.6B pixart 模型生成 512px 图片需 2 秒，而生成 4096px 图片，则需 3 分钟！在 V100 上生成 1028px 图片，1.5B HunyuanDiT 需 3 分 38 秒！

对于视频生成，单卡推理必然无法满足需求。视频生成问题延迟问题更严重，采用 DiT 架构的 Sora 推理成本让人高不可攀。OpenAI 的 CTO Murati 今年三月接受 WSJ 采访 [5] 时承认，Sora 对外开放使用的瓶颈之一正是推理开销。Factorial Funds 估算 Sora 的推理需求 [6]：一个 H100 GPU 一小时生成大约 5 分钟的视频，理想扩展下，需要 12xH100 实现 1 秒钟生成一秒钟视频。

越来越多的视频生成应用等待落地，为了满足它们在线服务满足实时性的需求，多 GPU 乃至多机部署 DiT 已势在必行。但是，我们却没有看到一个软件可以满足日益增多的 DiT 模型和实时性要求，更别说将 DiT 并行标准和实现统一起来。

大家目前都是基于 huggingface diffusers 库进行定制改造。正如 vLLM 出现之前，大家也都基于 huggingface transformers 库，增加并行逻辑去搭建推理流程，比如 deepspeed 和 accelerate 等，这些现在看来都是临时方案。势必需要一个性能好（满足实时性需求）的人人可用（易用性高）的 DiT 推理框架。

针对这个问题，我们将之前的 PipeFusion 项目升级成了通用的 DiT 推理引擎，命名为 xDiT。这个名字很像马斯克起的，不过我们是向 vLLM 致敬，这里的 X 是 Scalable 的意思。

论文标题：

xDiT: A Scalable Inference Engine for Diffusion Transformers (DiTs) on multi-GPU Clusters

项目地址：

https://github.com/xdit-project/xDiT

论文地址：

https://arxiv.org/abs/2405.14430

我们希望 xDiT 成为 DiT 推理领域的 vLLM。正如一年前，LLM 推理最关键的问题是如何做 Batching，而业界缺少有效的方案，所以 vLLM 携创新 Paged Attention+前人 Continous Batching 横空出世。

现在，DiT 推理最关键的问题是如何高效并行，业界的方案也比较匮乏，xDiT 携我们创新的 PipeFusion+ 前人的 Sequence Parallel 亮相。做一个比喻，vLLM=huggingface transformers+Batching，xDiT=huggingface diffusers+Parallel。

DiT 和 LLM 推理任务特点不一样，这导致二者软件设计哲学也有差异：

1. LLM 有 Prefill 和 Decode 两阶段，分别是计算密集和访存密集的；而 DiT 的计算和 Prefill 相似是计算密集的。所以 Batching、Paged Attention 等为了增加 Decode 阶段计算密度的设计在 DiT 中不需要，其实也没有 KVCache 不断 Append 增长的现象。

2. LLM 模型很大，而序列长度有限，比如 ShareGPT 的任务生成顶多几百个 token。而 DiT 正好反过来，模型不大，推理时序列长度很长，甚至 Million 级别。这导致张量并行这种 vLLM 中第一优先级的并行方法，在 DiT 中性能很低。

3. LLM 模型大多收敛到微调 Llama 架构，而 DiT 架构则呈现出较大的差异性。例如，SD3、Latte 等网络与 Meta DiT 论文中的 DiT Block 在算子设计上存在显著差异，且 DiT Block 之间的连接方式也各不相同，常采用类似 U-Net 的 U 型结构进行链接。因此，DiT 模型的代码相对 LLM 不够简洁，往往涉及好多文件和函数，对其进行模型层面的改造也更具挑战性。

正如 vLLM 让 transformers 具备 Batching 的能力，xDiT 让 diffusers 有并行的能力。考虑到 DiT 的特点，我们首先提出了一系列创新的并行方式，并可以相互混合。更重要的是，xDiT 提供了一套优雅的开发接口，针对性解决了 DiT 模型更改难度高的问题，这套开发接口尽可能可能复用 diffusers 的现有逻辑，开发者通过一些 wrapper 实现复杂的混合并行，实现高效地扩展方法。

xDiT Overview

作为可扩展的高性能推理 DiT 引擎，xDiT 有如下核心能力：

第一，DiT 主干网络混合并行：xDiT 支持任意混合四种基础并行策略方法。基础并行方法有：

Pipefusion Parallel [7]：TeraPipe 方式流水线并行方法，在弱互联，比如 pcie/以太网，的网络硬件上有优势，比如以太网互联的多机多卡情境下效率远高于其他方法。也是我们的创新工作。
Sequence Parallel：混合序列并行，使用 ring & ulysses sequence parallel 混合方式 USP [8]。
Data Parallel：在输入多个 prompt 或单个 prompt 生成多张图片时，在 image 间并行处理。
CFG Parallel, a.k.a Split Batch：在模型使用 classifier free guidance（CFG）时可以开启，并行度恒定为 2。

xDiT 中这四种并行方式可以以任何形式排列组合，混合策略在通信量和通信拓扑两方面可以取得最佳效果，使 xDiT 可达到近似线性的扩展性。

第二，Parallel VAE：针对扩散模型后处理的解码 VAE 模块在高分辨率图像生成时 OOM 问题，xDiT 实现了 Patch Parallel 版本的 VAE。

第三，简单灵活的开发接口：xDiT 轻松帮助用户支持新的 DiT 模型。它尽可能复用开源社区 diffuser，完成对模型推理的并行化改造。这套接口也是让 xDiT 从之前的不同并行方法集合变成一个灵活开发工具。通过 xDiT 接口，实测 15 分钟可以让 Pixart-Sigma 模型并行起来！

xDiT 本次的重构的核心，由一名本科生潘瑾哲 @Eigensystem 完成，瑾哲是华科超算队队长，架构功力超过了很多正式员工。关于项目更多细节，可以参考瑾哲文章 [9]。

▲ xDiT的方法

PipeFusion

PipeFusion 是一种 TeraPipe 式的序列流水并行。但和针对 LLM 设计 TeraPipe 有问题有显著不同，它利用了 Input Temporal Redundancy——即扩散步骤间输入与激活之间的数值相似性，这是扩散模型特有的特性，同样被应用于 DistriFusion 中。

Input Temporal Redundancy 和序列流水并行结合，避免了 LLM 模型中因果注意力带来的负载均衡问题。PipeFusion 可以显著减少了通信量，它和序列并行、张量并行、DistriFusion 对比都是通信带宽需求都是最低的，从而特别适合在以太网连接的多节点设置和 PCIe 连接的多 GPU 上使用。

PipeFusion 原理可以参考本人如下文章：

https://zhuanlan.zhihu.com/p/699612077

PipeFusion 和 LLM 序列并行对比可以参考本人如下文章：

https://zhuanlan.zhihu.com/p/706475158

USP：混合序列并行

序列并行我在之前写了太多文章，这里不在赘述，简单来说我们可以混合 ulysses 和 ring 两种并行，xDiT 的并行 args 也提供了 ring_degree 和 ulysses_degree 可以配置。

序列并行做大模型训练，你需要知道的六件事：

https://zhuanlan.zhihu.com/p/698031151

大模型训练之序列并行双雄：DeepSpeed Ulysses和Ring-Attention

CFG Parallel

Classifier-Free Guidance（CFG）是扩散模型领域的一个重要的技巧，可以提供更广泛的条件控制、减少训练负担、增强生成内容的质量和细节，以及提高模型的实用性和适应性。

对于一个输入 prompt，使用 CFG 需要同时进行 unconditional guide 和 text guide 的生成，相当于输入 DiT blocks 的 input latents batch_size = 2。CFG Parallel 分离两个 latents 分别进行计算，在每个 Diffusion Step forward 完成后、Scheduler 执行前 Allgather 一次 latent space 结果。它通信量远小于 Pipefusion 和 Sequence Parallel。

因此，使用 CFG 一定要优先使用 CFG Parallel。DistriFusion 也用了 CFG，叫 split batch，也是性能提升关键。

Hybrid Parallel

xDiT 设计目标是扩展 DiT 推理过程到超大规模，我们目标不止单机多卡，而是考虑了异构网络互联条件下多机多卡，比如以太网+PCIe。单一并行方式，比如 PipeFusion 或者 SP，很难同时做到这两点，不同并行方式混合在一起变得尤为重要。

如前文所述 xDiT 支持四种并行方式：PipeFusion、Sequence、Data 和 CFG Parallel。其中，Data 和 CFG Parallel 在图像间并行相对简单，而 PipeFusion 和 Sequence 在图像内部的不同 Patch 间并行则较为复杂。能让着两种并行方式的混合使用，正式 xDiT 核心创新点之一。

PipeFusion 和 USP 混合并不简单，因为 PipeFusion 使用上一个 Diffusion Step 过时的 Stale KV（Stale KV）参与计算，而传统 SP 都是只使用最新的 Fresh KV 参与计算。xDiT 针对这个新问题进行了特殊设计，保证混合并行执行结果准确，可以参考项目的设计文档 [10]。

我们在之前也是实现了 Tensor Parallel 和 DistriFusion，但发现他们都不足以作为基础并行方法。TP 通信量太大，DistriFusion KV 内存消耗太大，在 NVLink 上性能和 USP 近似，在 PCIe 弱互联又逊于 PipeFusion。不过，我们后续也考虑加入 DistriFusion 作为基础并行。

Parallel VAE

上述优化是对 DiT backbone 设计的，其实 VAE 模块在高清图片生成式也有问题。diffusers 采用的 stabilityai/sd-vae-ft-mse [11] 在处理高分辨率图像（A100上的 8192px）时会导致内存溢出（OOM）。这在 diffusers/issues/5924 [12] 中有详细记录，但是一直悬而未决。

我们借着 xDiT 机会彻底把这个问题解决了，为此开发了 DistVAE，并行高效处理高分辨率图像的解决方案：

https://github.com/xdit-project/DistVAE

我们的方法结合了两种关键策略：

1. 序列并行：我们将潜在空间中的特征图分割成多个 Patch，并在不同设备上进行序列并行 VAE 解码，就是序列并行。这种方法将中间激活所需的峰值内存减少到 1/N，其中 N 是使用的设备数量。

2. 分块输入处理：类似于 MIT-patch-conv [13]，我们将输入特征图分割成块，并依次送入卷积运算符。这种方法最小化了临时内存消耗。

通过整合这两种方法，我们显著扩展了 VAE 解码的能力。我们的实现成功处理了高达 10240px 的图像分辨率，与默认 VAE 实现相比，这有 11 倍提升。

未来工作

xDiT 项目我们未来计划做如下事情：

1. 更多模型支持。比如，最近出现的文生图模型 Flux，视频模型比如 CogVideo、Latte 等。我们也正在和 hugginface diffusers 团队建立合作，毕竟 xDiT 设计初衷也是通过添加少量代码让 diffusers 拥有并行能力。

2. ComfyUI 界面支持，让平民玩家可以通过 xDiT 把大的模型玩起来。比如，家里台式机配两张显卡，xDiT 可以有加速一倍。

3. xDiT 也可以有些调度层面的事情可以探索，比如 PipeFusion 的 Warmup 阶段分离。PipeFusion 流水并行 Warmup 会影响流水线，可以将预热阶段用 SP 单独部署，其余 diffusion step 用 PipeFusion。另外，KV 内存 Offloading 策略，LoRA 等方向也可以在分布式 DiT 推理中进行研究。

4. 和 Runtime 互补，比如 torch.compile，TensorRT，量化剪枝等。我们也正在和 OneDiff [14] 合作，进行编译优化方向的工作。xDiT 专注并行和存储管理，和计算加速兼容，可以组合成为最佳实践。

5. xDiT 也希望能汇聚更多 DiT Sys 方面最新研究，不仅仅是我们自己的方法。

希望对 xDiT 感兴趣的朋友与我们联系。我们也需要一些 researcher intern 一起做一些前沿方向。