--- title: 'KnapFormer: Efficient Load Balancing for DiT' url: https://www.emergentmind.com/topics/knapformer type: topic --- # KnapFormer: Efficient Load Balancing for DiT 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 [2508.06001]. ## 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 $s_j$ GPUs, $$ L_j = \sum_{i=1}^{N} \frac{w_i}{s_j} x_{ij} + \frac{\omega_j}{s_j} $$ where $x_{ij} = 1$ if sequence $i$ is assigned to bag $j$, and $0$ otherwise. The optimization seeks to minimize the deviation of $L_j$ from the mean per-GPU load $\bar{L}$: $$ \bar{L} = \frac{\sum_i w_i + \sum_j \omega_j}{G}, \qquad G = \sum_{j=1}^{M} s_j $$ The assignment $x_{ij}$ is determined by minimizing $\max_j |L_j - \bar{L}|$ 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 $l$, $$ \mathrm{FLOPs}(l) = 24\,l\,d^{2} + 4\,l^{2}d $$ Latency on H100 is dominated by the quadratic term, which is bandwidth-bound; the wall-clock cost is parameterized as $$ t = k \big(24\,l\,d^{2} + \gamma\,4\,l^{2}d \big) $$ $\gamma$ is empirically measured in the range $[0.38,\,0.49]$ for H100. The communication overhead $\omega_j$—primarily two all-to-alls per iteration for each bag—is modeled/measured per hardware, then converted into “equivalent computation” by dividing by $k$. Thus, only $\{w_i\}$ and $\{\omega_j\}$ 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 \}$ are first all-gathered within each data-parallel replica, $\mathcal O(N)$ ints per replica. b) **Workload and Assignment Computation:** On a designated rank/CPU, $w_i$ and $\omega_j$ are computed. A greedy first-fit–descending heuristic approximately solves the knapsack assignment $x_{ij}$ in $O(N\log N)$, 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$ across $s_j$ GPUs. d) **Sequence-Parallel Execution:** Bags with $s_j>1$ GPUs perform DeepSpeed-Ulysses `pre_attn`, an additional Ulysses all-to-all, local FlashAttention, and `post_attn`. Bags of size $1$ 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: ```python # (1) Gather metadata local_lens = [l1, l2, …, l_n] global_lens = all_gather(local_lens) # (2) Compute workloads on one rank if is_rank0: w = [24*li*d*d + γ*4*li*li*d for li in global_lens] ω = [comm_cost*2/k for each bag size] # two all2all per iter x = greedy_knapsack_solver(w, ω, bag_sizes) broadcast(x) # (3) Plan routing for each local sequence i: assigned_bag = arg j s.t. x[i,j]=1 s = bag_sizes[assigned_bag] split_offsets = chunk_positions(l_i, s) record (source_rank, target_ranks, offsets) # (4) Route chunks (one all_to_all per bag) bala_tokens = all_to_all(local_tokens, mapping) # (5) Forward/Backward for each block: (bala_q, bala_k, bala_v) = pre_attn(bala_tokens, seq_lens) bala_attn_out = FlashAttention_varlen(bala_q, bala_k, bala_v) bala_tokens = post_attn(bala_attn_out) # (6) Reverse routing out_chunks = all_to_all(bala_tokens, inverse_mapping) reassemble(out_chunks) → original order outputs ``` 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, $d_{\rm model}=3072$, 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 $1.00$ (i.e., $<1\%$ imbalance), lowers FBL to $3.44$s, increases TPS to $108$K, and HFU to $29.7\%$, realizing $2\times$–$3\times$ speedup and near-perfect balancing for token lengths spanning $10^{2}$ to $10^{4}$. 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 $\gamma$ is validated for token lengths $200$–$20\,000$, 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 [2508.06001]. Source: https://www.emergentmind.com/topics/knapformer