作者丨孙熙博
腾讯高级工程师，xDiT核心开发者
单卡 VS 并行版本在 CogVideoX-2b 上的推理效率
一、引言
今年 2 月 OpenAI Sora 发布，展示了令人惊艳的视频生成能力，随后各种文生视频产品如 ViDu、可灵、通义万相 等相继问世。尤其是字节的PixelDance千呼万唤始出来，更是让一些自媒体直呼Sora时代一去不复返了。这些产品对提示词（promt）的理解愈加精确，生成的视频也愈加逼真，不仅获得了用户的积极使用，也在商业化上取得成功。 可以预见，在不远的将来更多文生视频的新模型也将积极融入赛道。CogVideoX 作为目前文生视频领域最大的开源预训练模型，也开始受到广泛的欢迎。
DiT（Diffusion Transformer）是包含CogVideoX在内众多文生图与文生视频模型的核心网络结构。由于其优异的生成效果，近年来也逐渐取代 U-Net 成为内容生成的标准模型。DiT 属于一种扩散模型（Diffusion model），将初始图像或视频（通常为高斯噪声）输入以后，DiT 使用Transformer网络构架预测输入数据存在的噪声，并对其去噪，多次迭代即可实现图像或视频的生成。然而，由于模型本身的海量参数、较大视频尺寸、以及对生成视频时长需求的提升，推理过程中的高延迟成为CogVideoX等生成模型不可回避的问题，也使其无法应用在实时场景。
我们最近的工作 xDiT 旨在利用多 GPU 并行化 DiT 的推理过程，缩短推理延迟，冲击实时视频生成的场景需求。xDiT目前已支持Stable Diffusion 3、Flux、Pixart、HunyuanDiT、Latte 等众多模型，是高性能视频与图像生成的首选框架，不仅高效实现了 DiT 运算的通用并行策略，也针对每一种模型对网络结构的独特设计开发专有的并行方案。
GitHub链接：https://github.com/xdit-project/xDiT
二、认识CogVideoX
智谱开放平台最近推出了名为CogVideoX的视频生成模型，该模型采用了一种创新的Transformer架构，能够同时融合文本、时间和空间三个维度，实现不同模态信息的对齐和交互。
在使用CogVideoX生成视频的过程中，输入包括视频和文本提示，视频首先通过VAE模块从像素空间编码到潜在空间，然后在模型处理后返回像素空间；文本提示则被映射到连续空间，以便与视频数据一起处理。在潜在空间中，CogVideoX Transformer预测并去除视频噪声，经过去噪的视频会继续输入Transformer进行迭代。
CogVideoX 视频生成流程
CogVideoX Transformer 的细节流程展示如图。由于标准Transformer只能处理一维数据，而视频数据具有三个维度，即帧数、高度和宽度，因此在CogVideoX中，首先利用卷积层（Conv）对视频进行特征提取，以捕捉其空间信息。接着，将卷积层的输出展平 （Flatten）成一维的视频序列，以便适应Transformer的输入要求。为了保留视频序列中各个元素在原始视频中的位置信息，最后会添加位置编码。
CogVideoX Transformer 模型结果
在图右侧，CogVideoX Block 同时接收文本序列与视频序列，通过注意力机制将两者融合，其中符号⨁表示加法，⨂表示乘法。首先，归一化的文本和视频序列合并进行注意力计算，然后注意力的激活值被分开，经过LayerNorm和FeedForward层处理。TimeStamp则通过 Linear 层分别产生对文本和视频序列的 scale、shift 和 gate，其中前两者用于输入LayerNorm，而 gate 则用于与注意力的激活值相乘。视频序列在透过 30 个相同的CogVideoX Block后，通过 LayerNorm 和 Linear and Reshape 得到预测噪声。
由于 DiT 每轮迭代之间存在数据依赖，我们目前专注于在单次迭代内挖掘并行处理的潜力，以提高每轮噪声预测的计算效率。基于 xDiT框架，我们并行化 CogvideoX 的生成过程，实现了序列并行和 CFG 并行两种推理优化方法。
三、序列并行
序列并行使用多张 GPU 卡，每张卡保存全部的模型参数，但将视频序列拆分（Scatter）成多个碎片（Shard），将碎片分别传入 GPU 进行噪声预测，最后视频碎片组合（Gatter）成为完整视频，如下图所示。
采用的序列并行的 CogVideoX 视频生成流程
当输入数据从完整视频转变为视频碎片时，CogVideoX Transformer中的部分组件，如卷积层和注意力层，需要进行并行化改造，因为它们需要在多个序列元素间进行运算。相反，线性层、归一化层和FeedForward层由于仅进行元素级别的操作，可以直接应用于并行环境。
在xDiT框架中，已经积累了对常见网络模块的并行化最佳实践。例如，对于卷积模块，xDiT设计了多卡并行的卷积操作，保持输入输出与原始卷积模块对其，同时在模块内部处理通信细节，并利用缓存策略减少内存重分配的开销。因此，用户可以直接将原始网络的卷积层替换为xDiT的并行版本，无需深入理解多卡间的通信细节，以最小的时间成本实现高效并行卷积运算。
随着长序列生成的需求增长，注意力机制的序列并行也变得至关重要。现有的并行方法包括Ulysses Attention和Ring Attention，前者通过运算之前的 All2All 原语完成数据交换，而后者在运算过程中借助Peer-to-Peer通信传输视频碎片。这两种注意力模块的并行策略都已整合到 xDiT 中。xDiT 还实现了 USP（Unified Sequence Parallelism），同时结合了这两种策略，支持灵活的并行模式组合。
然而，上述方法只针对单一输入序列做注意力运算，而CogVideoX的Full Attention计算涉及文本和视频序列的组合，这为并行化带来了新的挑战。通过分析CogVideoX模型各层权重的维度，我们发现文本序列长度固定为226，而视频序列长度与生成视频的帧数相关，通常很长（例如17550）。鉴于两者巨大的长度差异，我们为每台GPU 分配一个视频碎片，但将整个文本序列存在所有 GPU 上。这样，我们可以完全复用 xDiT的视频拆分逻辑，只需对现有的并行注意力运算进行微小调整以支持文本和视频之间的Full Attention：在传输时，由于每台GPU 存储了全部文本序列，All2All 和Peer-to-Peer传输中涉及的文本部分将被省略；在计算时，当GPU处理视频序列的前端部分时，它同时处理对文本的注意力运算。
四、CFG并行
CogVideoX 还采用了 Classifier-Free Guidance（CFG）技术来提升生成内容的质量和细节。
同时采用的序列并行和 CFG 并行的 CogVideoX 视频生成流程
每轮迭代中，模型不仅将视频和文本输入到CogVideoX Transformer，还会生成一个空文本，并与相同的视频一起传递给另一个Transformer副本。在每轮迭代结束后，CogVideoX会融合两个Transformer的输出来计算噪声。CFG并行化不需要进一步分割视频序列，通信开销较小，因此实现起来相对简单。在 xDiT框架中，我们为每个视频碎片在两张GPU卡上分别部署网络模型。在每轮迭代中，这两张GPU独立执行计算，然后在迭代后交换输出数据以获得最终的噪声。这种设计使得CFG并行化在不增加复杂性的同时，能够有效地利用多GPU资源，提高生成效率。
五、性能分析
上述两种并行策略，即序列并行和Classifier-Free Guidance（CFG）并行，已经被整合到xDiT项目中。在配备了L40（PCIe）GPU的系统上，我们对比了使用diffusers库的单卡CogVideoX推理与我们设计的并行版本在生成49帧（6秒）720x480分辨率视频时的性能差异。
如图所示，在基础模型CogVideoX-2b上，无论是使用Ulysses Attention、Ring Attention还是CFG并行，推理延迟都有所减少。其中，CFG并行由于通信量较小，表现出比其他两种技术更高的性能。通过结合序列并行和CFG并行，我们进一步提高了推理效率。随着并行度的增加，推理延迟持续降低。在最佳配置下，xDiT相比于单卡推理方法可以实现3.53倍的加速，每次迭代只需0.6秒。由于CogVideoX默认的迭代次数为50，总计30秒即可端到端生成长达6秒视频。
对于更复杂的CogVideoX-5b模型，它拥有更多的参数以提升视频生成的质量和视觉效果，但计算成本也显著增加。尽管如此，所有方法在该模型上依然保持了与CogVideoX-2b模型相似的性能趋势，但并行版本的加速比得到了进一步提升。在与单卡版本的比较中，xDiT可以实现高达3.91倍的推理速度提升，端到端生成视频的时间仅需80余秒。
六、总结
CogVideoX在视频生成方面的出色表现，赢得了业界的广泛赞誉，进一步证明了DiT模型在文本生成图像和视频领域的巨大潜力。作为DiT模型的多卡并行推理框架，xDiT成功地适应并优化了现有的图像和视频生成流程，同时对多种模型组件进行了并行化处理。随着新模型的发布，它们可以轻松地整合到xDiT中，从而实现显著的性能提升。