Blog Tags GitHub
Back to blog
diffusion-model system-optimization parallelism-sp parallelism-pp parallelism-tp

xDiT: an Inference Engine for Diffusion Transformers (DiTs) with Massive Parallelism

面向扩散 Transformer 的大规模并行推理引擎，混合多种并行策略实现 16 GPU 扩展

xDiT: an Inference Engine for Diffusion Transformers (DiTs) with Massive Parallelism

一、论文概述

项目内容
标题xDiT: an Inference Engine for Diffusion Transformers (DiTs) with Massive Parallelism
作者Jiarui Fang, Jinzhe Pan, Xibo Sun, Aoyu Li, Jiannan Wang
机构Tencent
论文https://arxiv.org/abs/2411.01738
代码https://github.com/xdit-project/xDiT
发布2024-11-04

二、核心思想

问题定义

扩散 Transformer (DiT) 在图像和视频生成中展现出卓越能力，但生成高质量内容需要更长的序列长度，导致注意力机制的计算量呈指数级增长，推理延迟严重。例如：

  • Flux.1 生成 1024px 图像需要 262K token 序列
  • 4096px 分辨率图像需要 420 万 token
  • CogVideoX 生成 6 秒视频需要超过 100 万 token

实时部署 DiT 必然需要多设备并行推理，但单一并行方法在大规模场景下扩展性差。

三大挑战

  1. 硬件网络多样性：集体通信方法仅适合 NVLink 等高带宽互连，P2P 方法更适合 PCIe/Ethernet
  2. 模型架构多样性：不像 LLM 有统一架构，DiT 在指令注入方式、块连接方式、算子布局等方面差异很大
  3. 长序列和大模型的内存需求：当前 SOTA 图像生成模型达 120 亿参数

解决方案概述

xDiT 是一个综合并行推理引擎，灵活混合三种并行策略：

  1. Sequence Parallelism (SP)：图内并行，沿序列维度分割
  2. PipeFusion：图内并行，Patch 级流水线并行（本文提出的新方法）
  3. CFG Parallelism：图间并行，分离条件/无条件推理

三、技术架构

并行策略详解

1. Sequence Parallelism (SP)

将注意力计算沿序列维度分配到多个 GPU：

  • 每个 GPU 只计算部分 token 的注意力
  • 需要高带宽互连（NVLink）进行 AllReduce 通信
  • 适合高带宽网络环境

2. PipeFusion: Patch-level Pipeline Parallelism（本文核心贡献）

PipeFusion 将输入潜在图像分割成 patch，将 DiT 模型分割成层：

分区策略：

  • 模型沿数据流分割，将连续层分配给不同 GPU
  • 输入图像分成 M 个不重叠的 patch
  • 每个 GPU 用分配的层并行处理一个 patch

流水线工作流：

  • 假设 N=4 GPU，M 个 patch
  • GPU 0 处理 patch 0 的第 0 层，然后处理 patch 1 的第 0 层…
  • GPU 1 等待 GPU 0 完成 patch 0 的第 0 层后，处理 patch 0 的第 1 层…
  • 形成流水线，各 GPU 持续工作

关键优势：

  • 通信效率高：使用 P2P 通信，适合 PCIe/Ethernet
  • 内存效率高：每个 GPU 只存储部分模型层
  • 计算效率高：DiT 由相同的 Transformer 块组成，易于均匀分区

3. CFG Parallelism（图间并行）

Classifier-Free Guidance 需要同时生成无条件引导和文本引导，等效于 batch size=2：

  • 将两个 latent 分离到不同 GPU 计算
  • 每个扩散步骤完成后，在 scheduler 执行前进行 Allgather
  • 通信开销远小于 PipeFusion 和 SP

4. 混合并行

xDiT 灵活组合以上策略：

  • PipeFusion + CFG：适合 PCIe/Ethernet
  • SP + CFG：适合 NVLink
  • PipeFusion + SP + CFG：大规模场景
  • PipeFusion + SP：单图场景（无 CFG）

并行 VAE

VAE 解码器也面临高分辨率下的内存和计算瓶颈，xDiT 对 VAE 也实现了并行化。

四、核心创新

创新点说明理论/实验依据
PipeFusionPatch 级流水线并行，P2P 通信适合 PCIe/Ethernet，通信效率高于 SP
混合并行框架灵活组合 SP、PipeFusion、CFG适应不同硬件网络和模型架构
首个 16 GPU DiT 扩展首次将 DiT 推理扩展到 16 GPU5 个 SOTA DiT 模型验证
Ethernet 集群支持首次展示 DiT 在 Ethernet GPU 集群的扩展性2×8 L40 节点通过 Ethernet 互连
并行 VAEVAE 解码器的并行化高分辨率下的内存和计算优化

五、实验结果

测试环境

  • L40 集群：2×8 NVIDIA L40 GPU (PCIe)，Ethernet 互连
  • A100 集群：1×8 NVIDIA A100 GPU，NVLink 互连

图像生成性能（L40 集群）

Pixart 模型：

  • 4096px 任务：16 GPU 实现 13.29× 加速，延迟从 245s 降至 17s
  • 最佳混合并行策略随图像大小变化

Flux.1 模型：

  • 需要 120 亿参数，内存需求大
  • PipeFusion 显著减少内存开销

Stable-Diffusion3 模型：

  • 1024px 和 2048px 任务：pipefusion=8 + cfg=2 策略延迟最低
  • pipefusion=4 + ulysses=2 + cfg=2 次优

图像生成性能（A100 集群）

  • NVLink 互连下 SP 性能更好
  • 混合并行仍优于单一方法

关键发现

  1. PipeFusion 通信效率高：比 SP 和 TP 通信开销更低
  2. PipeFusion 内存效率高：在大模型（如 Flux.1）上显著减少内存
  3. 混合并行最优：单一方法无法适应所有场景
  4. Ethernet 可行：首次证明 DiT 可在 Ethernet 集群上扩展

质量评估

  • 并行化不影响生成质量（FID、IS 等指标无显著变化）

六、相关工作

方法类型局限性
Sequence Parallelism图内并行需要高带宽互连
Tensor Parallelism图内并行需要高带宽互连
DistriFusion图内并行P2P 通信，但扩展性有限
PipeFusion（本文）图内并行P2P 通信，流水线效率高
CFG Parallelism图间并行仅 batch 维度

七、总结

核心贡献

  1. 提出 PipeFusion：Patch 级流水线并行方法，P2P 通信适合 PCIe/Ethernet
  2. 混合并行框架：灵活组合 SP、PipeFusion、CFG 三种策略
  3. xDiT 系统：首个将 DiT 推理扩展到 16 GPU 的系统
  4. Ethernet 集群验证：首次展示 DiT 在 Ethernet 互连 GPU 集群的扩展性
  5. 5 个 SOTA 模型验证：Pixart、SD3、Flux.1、HunyuanDiT、CogVideoX

技术影响

  • 实用价值：xDiT 已开源，可直接用于生产部署
  • 硬件适应性：支持 PCIe/Ethernet/NVLink 等多种互连
  • 可扩展性：将 DiT 推理扩展到 16 GPU

局限性

  • 主要关注推理，未涉及训练
  • PipeFusion 的流水线气泡在 GPU 数量增多时可能增加
  • VAE 并行化效果受限于 VAE 架构

八、参考资源