KnapFormer: Efficient Load Balancing for DiT

• KnapFormer is an online load balancing framework that uses a knapsack-based global token redistribution scheme to balance workload across GPUs during DiT training.
• It integrates dynamic sequence parallelism and a semi-empirical compute cost model to minimize straggler delays and reduce communication overhead.
• Empirical results on the FLUX DiT model using 32×H100 GPUs show up to 3× speedup with near-perfect balancing, demonstrating significant improvements in hardware utilization.

KnapFormer is an online load balancing framework designed to enable efficient and scalable distributed training of Diffusion Transformers (DiT) under high input heterogeneity. By integrating a knapsack-based global token redistribution scheme and dynamic sequence parallelism, KnapFormer addresses the severe workload imbalance that arises in DiT models due to variable-length textual inputs and variable token counts from mixed-resolution or image-video data. The method minimizes straggler effects and maximizes hardware utilization by minimizing variance in per-GPU token-processing load, while accounting for both computational and communication overheads. KnapFormer supports state-of-the-art DiT models such as FLUX and is open-sourced for community adoption (Zhang et al., 8 Aug 2025).

1. Motivation and Problem Formulation

Distributed DiT training encounters significant token imbalance caused by heterogeneous input—particularly multi-modal, mixed-resolution, and variable-length samples. Standard data or sequence parallelism leads to uneven load distribution, manifesting as straggler delays and suboptimal hardware efficiency. KnapFormer introduces a principled solution by globally redistributing sequence tokens across participating GPUs to approach perfect per-GPU workload balancing, while minimizing communication overhead introduced by sequence-parallel computation.

The assignment problem is posed as a multi-knapsack optimization over $N$ input sequences and $M$ compute "bags" (each comprising $s_j$ GPUs) within a training replica. Let $l_i$ denote the length of sequence $i$, $d$ the model dimension, and $\gamma$ the fitted bandwidth-correction factor. The estimated compute cost for each sequence is

$w_i = 24\,l_i\,d^{2} + \gamma \cdot 4\,l_i^{2}d$

The effective per-GPU workload in bag $j$ includes both computation and the equivalent workload cost of sequence-parallel all-to-all collectives, $\omega_j$, evenly distributed across $M$0 GPUs,

$M$1

where $M$2 if sequence $M$3 is assigned to bag $M$4, and $M$5 otherwise. The optimization seeks to minimize the deviation of $M$6 from the mean per-GPU load $M$7:

$M$8

The assignment $M$9 is determined by minimizing $s_j$0 or the load variance.

2. Semi-Empirical Model of Workload and Communication

Accurate load estimation is critical for balancing. KnapFormer adopts a semi-empirical model of per-sequence compute cost based on transformer FLOP analysis and hardware profiling. For a sequence of length $s_j$1,

$s_j$2

Latency on H100 is dominated by the quadratic term, which is bandwidth-bound; the wall-clock cost is parameterized as

$s_j$3

$s_j$4 is empirically measured in the range $s_j$5 for H100. The communication overhead $s_j$6—primarily two all-to-alls per iteration for each bag—is modeled/measured per hardware, then converted into “equivalent computation” by dividing by $s_j$7. Thus, only $s_j$8 and $s_j$9 are required for solver input.

3. Algorithmic Design and Integration with DeepSpeed-Ulysses

At each training step, KnapFormer executes an online dataflow as follows:

a) Metadata Gathering: All sequence lengths $l_i$0 are first all-gathered within each data-parallel replica, $l_i$1 ints per replica.

b) Workload and Assignment Computation: On a designated rank/CPU, $l_i$2 and $l_i$3 are computed. A greedy first-fit–descending heuristic approximately solves the knapsack assignment $l_i$4 in $l_i$5, incorporating communication overheads.

c) Routing Plan and Chunk Distribution: Assignment mapping is broadcast globally. Each GPU chunks its assigned sequences based on bag size, distributing tokens via a single intra-bag all-to-all (PyTorch all_to_all), splitting each $l_i$6 across $l_i$7 GPUs.

d) Sequence-Parallel Execution: Bags with $l_i$8 GPUs perform DeepSpeed-Ulysses pre_attn, an additional Ulysses all-to-all, local FlashAttention, and post_attn. Bags of size $l_i$9 bypass collectives.

e) Reverse Routing: After forward and backward passes, an inverse mapping all-to-all restores original sequence order for loss and optimizer computation.

The end-to-end communication requirement is restricted to two small intra-bag all-to-alls per iteration per bag, which sharply contrasts the high overhead of naive per-block rebalancing.

4. Workflow Summary and Pseudocode

The KnapFormer scheduling pipeline at each training iteration is as follows:

$d$3

The “greedy_knapsack_solver” uses a three-phase approach to produce near-optimal assignments efficiently.

5. Empirical Evaluation and Results

Evaluation was performed on the FLUX DiT model (12B parameters, $i$0, 19 DoubleStream + 38 SingleStream blocks) with 32 × H100 GPUs, using three data regimes:

• Low-resolution image (256×256 + text)
• Mixed-resolution image (256–2048 px buckets)
• Joint image+video (keyframes, short/long clips)

Four homogeneous bag topologies were tested: g1n32, g2n16, g4n8, g8n4. Key metrics are:

• Workload Imbalance Ratio (WIR)
• Forward+Backward Latency (FBL)
• Throughput (Tokens/sec)
• Hardware-FLOP Utilization (HFU)

Principal findings are:

Regime (no KnapFormer)	WIR	FBL (s)	TPS (K/s)	HFU (%)
Mixed/joint data	up to 28	≈ 8	≈ 46	≈ 13

KnapFormer (g8n4) reduces WIR to $i$1 (i.e., $i$2 imbalance), lowers FBL to $i$3s, increases TPS to $i$4K, and HFU to $i$5, realizing $i$6–$i$7 speedup and near-perfect balancing for token lengths spanning $i$8 to $i$9. For low-variation corpora, g1n32 suffices, but for highly heterogeneous corpora, g8n4 (8-way intra-node) is optimal.

6. Context and Implications

KnapFormer demonstrates that global knapsack-based token redistribution, coupled with explicit modeling of compute and communication costs, can substantially reduce hardware underutilization and eliminate workload stragglers. By leveraging real-time all-gather metadata exchange and efficient parallel heuristics, it achieves near-ideal balancing at minimal communication cost. The semi-empirical model’s accuracy for $d$0 is validated for token lengths $d$1–$d$2, supporting broad dataset and sequence heterogeneity.

A plausible implication is that similar formulations can be adapted to other distributed sequence-processing architectures exhibiting workload variability, especially when leveraging complex parallelization strategies such as those in DeepSpeed-Ulysses. This suggests that explicit handling of both token imbalance and sequence-parallel overhead is critical for scaling DiT and related models to real-world, heterogeneous corpora (Zhang et al., 8 Aug 2025).