SwiftFusion: Scalable Sequence Parallelism for Distributed Inference of Diffusion Transformers on GPUs
本文介绍了 StreamFusion，这是一种面向扩散 Transformer 的拓扑感知分布式推理引擎，它利用新颖的环面注意力机制和单向通信来克服延迟与同步瓶颈，相比最先进的方法实现了高达 1.77 倍的加速。
原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。 阅读完整免责声明
想象一下，你正试图用一份需要成千上万个小步骤的食谱，烘焙一个巨大的多层蛋糕（即一张高质量图像或视频）。在人工智能领域，这个“蛋糕”是由扩散变换器（Diffusion Transformer, DiT） 生成的。
问题在于，随着蛋糕变得更大（分辨率更高或视频更长），单个厨房烤箱（单块 GPU）无法承受如此大的工作量。它变得太慢，而且原料（数据）占用的空间也太大。因此，我们雇佣了一个烤箱团队（多块 GPU）协同工作。
然而，仅仅把烤箱并排摆放是不够的。它们需要不断地互相传递原料。如果它们花在传递碗上的时间比烘焙的时间还多，整个过程就会变慢。这正是SwiftFusion所要解决的问题。
以下是该论文如何用三个主要概念来解释他们的解决方案：
1. “高速公路与土路”问题（拓扑感知调度）
想象你的烤箱团队被分成两组：
- A 组：位于同一厨房内的烤箱，通过超高速、宽大的传送带连接（机内网络）。
- B 组：位于不同建筑内的烤箱，通过较慢、狭窄的土路连接（机间网络）。
以前的方法试图用慢土路来承担繁重任务，用快传送带处理轻量任务。这就像试图把一张巨大的沙发搬过一条狭窄的小巷——导致了交通堵塞。
SwiftFusion 的解决方案：他们颠覆了这一策略。
- 他们利用快速传送带来处理繁重且杂乱的原料传递（Ring Attention）。
- 他们仅利用慢速土路来进行有组织的大件货箱运输（Ulysses Attention）。
通过将任务与正确的“道路”相匹配，他们避免了堵塞慢速连接。
2. “流水线”技巧（环面注意力，Torus Attention）
在旧方法中，烤箱必须排队等待。烤箱 1 把碗传给烤箱 2，等待烤箱 2 完成，然后烤箱 2 再传给烤箱 3。这造成了大量的“死时间”，烤箱只是在那里等待。
SwiftFusion 的解决方案：他们将其转变为繁忙的流水线。
他们不再等待整个碗到达后再开始工作，而是将碗分成小块。
- 一旦烤箱 1 从邻居那里收到第一块原料，它就立即开始烘焙这一块。
- 在烘焙第一块的同时，它同时接收第二块。
- 当第二块到达时，烤箱已准备好立即烘焙它。
这被称为环面注意力（Torus Attention）。这就像一位厨师，在送货卡车仍在卸载其余农产品时，就开始切蔬菜。他们将“等待”（通信）与“工作”（计算）重叠，从而不浪费任何时间。
3. “自助式”配送（单边通信）
在旧系统中，传递原料需要“握手”。烤箱 1 会大喊：“我要发送这个了！”，而烤箱 2 必须回喊：“我准备好接收了！”这种不断的喊叫和等待造成了大量的噪音和延迟（同步开销）。
SwiftFusion 的解决方案：他们使用一个名为 NVSHMEM 的特殊库，切换到了“自助式”模式。
- 现在，烤箱 1 可以直接将一盘原料滑到烤箱 2 的柜台上，而无需事先请求许可。
- 烤箱 2 可以在准备好时随时取走托盘。
- 这消除了不断喊叫和等待的需求，使整个厨房运行得更加顺畅。
结果
该论文在生成高分辨率图像和长视频时测试了这个新系统。他们发现，通过使用这三个技巧：
- SwiftFusion 的平均速度比当前最佳方法快 1.35 倍。
- 在最佳情况下，速度快 1.77 倍。
- 它不需要更多的内存（原料），只需要一种更聪明的移动方式。
简而言之，SwiftFusion 通过更好地组织“厨房”、让烤箱在等待时也能工作，并消除它们之间不必要的“握手”，从而加快了 AI 图像和视频生成的速度。
您所在领域的论文太多了？
获取与您研究关键词匹配的最新论文每日摘要——附技术摘要，使用您的语言。