KnapFormer动态调度扩散模型训练负载，提升多模态训练效率
KnapFormer: An Online Load Balancer for Efficient Diffusion Transformers Training
- 通过求解全局背包问题重分配序列数据，平衡各GPU工作量
- 实测负载差异低于1%，训练速度提升2至3倍
- 适合大规模图像视频联合训练的扩散模型研究者使用
我们提出KnapFormer，一种高效通用的分布式训练框架，用于扩散变换器（DiT）的负载均衡与序列并行结合。其核心洞察是：序列并行与跨GPU的显著分块不均衡之间存在强协同效应。该不均衡源于可变长度文本输入及混合分辨率、图文联合训练中视觉标记数量差异。KnapFormer首先在平衡组内汇聚各节点的序列长度元数据，再求解全局背包问题，目标是最小化每张GPU的总工作量方差，同时考虑序列并行的影响。通过在负载决策中集成基于DeepSpeed-Ulysees的序列并行机制，并采用简单半经验工作量模型，KnapFormer实现极低通信开销，在序列长度从数百到数万的现实训练任务中，负载偏差小于1%。它有效消除慢节点效应，在混合分辨率和图文联合数据集上训练FLUX等前沿扩散模型时，获得2至3倍加速。代码已开源：https://github.com/Kai-46/KnapFormer/
原文摘要 · Abstract (English)
We present KnapFormer, an efficient and versatile framework to combine workload balancing and sequence parallelism in distributed training of Diffusion Transformers (DiT). KnapFormer builds on the insight that strong synergy exists between sequence parallelism and the need to address the significant token imbalance across ranks. This imbalance arises from variable-length text inputs and varying visual token counts in mixed-resolution and image-video joint training. KnapFormer redistributes tokens by first gathering sequence length metadata across all ranks in a balancing group and solving a global knapsack problem. The solver aims to minimize the variances of total workload per-GPU, while accounting for the effect of sequence parallelism. By integrating DeepSpeed-Ulysees-based sequence parallelism in the load-balancing decision process and utilizing a simple semi-empirical workload model, KnapFormers achieves minimal communication overhead and less than 1% workload discrepancy in real-world training workloads with sequence length varying from a few hundred to tens of thousands. It eliminates straggler effects and achieves 2x to 3x speedup when training state-of-the-art diffusion models like FLUX on mixed-resolution and image-video joint data corpora. We open-source the KnapFormer implementation at https://github.com/Kai-46/KnapFormer/
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。