通过动态稀疏性加速视频扩散模型训练，提升3倍以上吞吐量。
DSV: Exploiting Dynamic Sparsity to Accelerate Large-Scale Video DiT Training
- 利用查询与键的低秩近似捕捉动态注意力稀疏模式。
- 在128张GPU上实现520k序列长度训练，吞吐量提升3.02倍无质量损失。
- 适合大规模视频生成任务，尤其适用于高分辨率长视频训练。
扩散变压器（DiT）在生成高质量视频方面表现出色，但3D全注意力机制的二次复杂度成为扩展训练的瓶颈，尤其在高分辨率、长视频场景下，该部分可占用高达95%的处理时间，并需专用上下文并行。本文提出DSV，通过利用我们实证观察到的动态注意力稀疏性来加速视频DiT训练。DSV采用两阶段算法，基于低秩近似捕获原始查询与键的动态稀疏模式，并使用自定义核函数高效识别关键的键值对并计算稀疏注意力。为适应新的稀疏维度，DSV引入混合稀疏感知上下文并行策略，重新平衡因稀疏异质性导致的注意力头与块间负载不均问题。DSV在不损失质量的前提下，将训练吞吐量最高提升3.02倍，支持扩展至128张GPU和520k token长度。
原文摘要 · Abstract (English)
Diffusion Transformers (DiTs) have shown remarkable performance in generating high-quality videos. However, the quadratic complexity of 3D full attention remains a bottleneck in scaling DiT training, especially with high-definition, lengthy videos, where it can consume up to 95% of processing time and demand specialized context parallelism. This paper introduces DSV to accelerate video DiT training by leveraging the dynamic attention sparsity we empirically observe. DSV uses a two-stage algorithm to capture the dynamic sparsity patterns via low-rank based approximation of the original query and key. It employs custom kernels to efficiently identify critical key-value pairs and compute the sparse attention. To accommodate the new sparsity dimension, DSV adopts a hybrid sparsity-aware context parallelism that re-balances the skewed workload across attention heads and blocks due to sparsity heterogeneity. DSV achieves up to 3.02x higher training throughput, scaling to 128 GPUs and 520k token lengths, without quality loss.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。