解决视频扩散模型训练中序列长度差异导致的算力浪费问题
AdaptiveLoad: Towards Efficient Video Diffusion Transformer Training
- 动态调整数据加载策略，同时控制显存和计算负载
- 训练不平衡率降至18.9%，显存利用率提升22.7%
- 适合大规模视频生成与世界模型训练的开发者
在视频生成模型（尤其是世界模型）中，训练大规模视频扩散变换器（如DiT和MMDiT）面临显著计算挑战，源于混合模式数据集中序列长度的极端差异。现有基于分桶的数据加载策略通常依赖‘等长令牌’约束，未能考虑自注意力机制的二次复杂度，导致严重负载不均和GPU资源闲置。本文提出 extit{AdaptiveLoad}，一个集成优化框架，包含两个核心组件：(1) 双约束自适应负载均衡系统，通过同时限制内存消耗和计算负载（$B \times S^p \le M_{\text{comp}}$）消除长序列瓶颈；(2) 融合层归一化-调制的CUDA内核，采用D-tile合并归约策略提升吞吐量并缓解内存压力。在Wan 2.1世界模型上的实验表明，该方法将计算不平衡率从39%降低至18.9%，峰值显存利用率提升22.7%，整体训练吞吐量提高27.2%。
原文摘要 · Abstract (English)
In video generation models, particularly world models, training large-scale video diffusion Transformers (such as DiT and MMDiT) poses significant computational challenges due to the extreme variance in sequence lengths within mixed-mode datasets. Existing bucket-based data loading strategies typically rely on "equal token length" constraints. This approach fails to account for the quadratic complexity of self-attention mechanisms, leading to severe load imbalance and underutilization of GPU resources. This paper proposes \textit{AdaptiveLoad}, an integrated optimization framework consisting of two core components: (1) A dual-constraint adaptive load balancing system, which eliminates long-sequence bottlenecks by simultaneously limiting memory consumption and computational load ($B \times S^p \le M_{\text{comp}}$); (2) A fused LayerNorm-Modulate CUDA kernel, which utilizes a D-tile coalesced reduction strategy to increase throughput and alleviate memory pressure. Experimental results on the Wan 2.1 world model demonstrate that our method reduces the computational imbalance rate from 39\% to 18.9\%, improves peak VRAM utilization efficiency by 22.7\%, and achieves an overall training throughput increase of 27.2\%.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。