KnapFormer: An Online Load Balancer for Efficient Diffusion Transformers Training

TL;DR

KnapFormer is an online load balancer for Diffusion Transformer (DiT) training under highly heterogeneous data — mixed images, keyframes, and videos at varying resolutions and frame rates. It logically groups GPUs into “compute bags” spanning one or more devices and dynamically re-routes sequence chunks across bags using a knapsack-style allocator, on top of Ulysses-style sequence parallelism, so per-step work is balanced across ranks. Ships as a minimally-intrusive integration with MMDiT-style architectures (Flux), with a workload estimator, a training simulator, and a one-click experiment launcher; targets multi-node FSDP setups at >=32 GPUs. Authors are from Adobe Research; the repo may relocate to adobe-research and current Kai-46 copy will mirror.

Key claims

The training bottleneck targeted is per-step imbalance, not per-GPU FLOPs: with heterogeneous sequence lengths (low-res image to high-res video) under naive sharding, some ranks finish forward+backward seconds before others and stall at the AllReduce/AllGather boundary [README “Description”].
The proposed primitive is a compute bag — one or more GPUs treated as a single logical unit for sequence-parallel attention — with a planner that re-routes sequence chunks across bags to equalize work per step, formulated as a knapsack-style assignment [README “Description”; visualization section].
A workload estimator benchmarks DiT and MMDiT (Flux) forward+backward cost as a function of d_model, d_head, causal, use_flux, n_ds_layers, n_ss_layers, with theoretical-vs-actual plots, so the routing planner consumes calibrated cost models rather than uniform-token assumptions [README “Workload Estimator”].
Integration with MMDiT is scoped to three small surfaces: a forward-pass routing/reverse-routing pair (simulator_model.py lines 75–146), pre/post-attention ops (transformer_utils.py lines 396–429), and FSDP-compatible conditional execution inside double-stream blocks (mmdit_utils.py lines 308–319) [README “Key Integration Points”].
Reported metrics from the simulator are forward+backward latency, total throughput, HFU (Hardware Flop Utilization), and load-imbalance ratio — with comparative runs with and without sequence balancing as the primary ablation [README “Training Simulator”].
Pre-configured Flux experiment grid requires at least 32 GPUs; smaller setups need a custom --balancer_config argument [README “Experiment Launcher”].
Apache 2.0 licensed; PDF tech report shipped in-repo as assets/paper.pdf [README “License” + “Citation”].

Method

The load balancer sits on top of Ulysses-style sequence parallelism. In Ulysses each attention layer’s tokens are sharded along the sequence dimension across a chosen group of GPUs, with an all-to-all exchange swapping the shard axis between sequence-parallel and head-parallel for the attention compute. KnapFormer’s contribution is to make that “chosen group” — what the codebase calls a compute bag — dynamic and heterogeneous across the batch.

A training batch contains sequences of very different lengths (a 1080p video frame stack is many times longer than a 256x256 image). The planner reads each sequence’s length plus the calibrated cost model from the workload estimator, then solves a knapsack-style routing problem: assign sequence chunks to compute bags so that every bag’s per-step compute is approximately equal, and bag sizes (GPUs per bag) are chosen to fit the largest chunks. The result is a routing plan — a JSON-serializable summary returned by balancer.get_routing_plan_summary() — that the training loop consults each step.

Concretely the MMDiT integration looks like this. On entry to a double-stream block, the router collects per-rank token shards and either redistributes them across bags (when the planner says a different bag assignment is cheaper) or leaves them in place. Pre- and post-attention ops adjust the active sequence-parallel group for the FlashAttention call. Reverse routing at block exit puts tokens back where they came from before residual addition. Double-stream blocks under FSDP gate text-token vs image-token paths with an FSDP-compatible conditional so that ranks holding zero tokens for a stream don’t break the all-reduce contract.

There is also a Manim-based visualizer (route_visualization.py) that animates the routing plan — useful for sanity-checking that a custom data mix doesn’t produce pathological assignments.

Results

This is a research codebase, not a paper page, so quantitative numbers are not in the README. What’s published:

The simulator measures forward+backward latency, total throughput, HFU, and load-imbalance ratio under matched --data_codes (e.g. g8b32i256f1s0) and --balancer_config (e.g. g1n8, g2n16) settings.
The experiment launcher generates aggregated reports comparing balanced vs unbalanced configurations across multiple Flux balancer setups.
A balancing animation (assets/animation.gif) shows GPU bags absorbing variable-length sequence chunks; visually demonstrates that bag size adapts to chunk size rather than the other way around.
An accompanying tech report PDF is in the repo (assets/paper.pdf) but the README does not quote numbers from it.

For an actual throughput table or speedup factor a reader needs to either run the included Flux experiment grid (≥32 GPUs) or fetch the PDF.

Why it’s interesting

KnapFormer is the DiT-side analogue of the LLM-training-systems papers the wiki has been tracking, and it isolates the one structural feature that distinguishes DiT pretraining from LLM pretraining: tokens-per-sample varies by orders of magnitude across a single batch. LLM training assumes roughly uniform sequence lengths (or pads to one), so veScale-FSDP: Flexible and High-Performance FSDP at Scale focuses on optimizer-state and FSDP buffer shapes, and SpectraX — True MPMD Pipeline Parallelism for JAX focuses on pipeline-stage heterogeneity. KnapFormer assumes the data is heterogeneous and the parallelism plan must vary per-step — a different bottleneck that vanilla FSDP, TP, or Ulysses-SP all leave on the table.

It also slots into the wiki’s diffusion-training-efficiency thread (Diffusion training efficiency) as a systems-side complement to the algorithmic-side recipes (REPA, SRA, Self-Flow, UL). Those papers reduce iterations to target FID; KnapFormer reduces wall-clock per iteration at given iteration count by recovering the ~10–40% of GPU time that imbalanced bags spend stalled. The two levers compose: a Self-Flow run on heterogeneous video-and-image data would benefit from both. Worth noting that Scaling and Optimizing Frontier Model Training (Fireworks AI)‘s “Context” parallelism dimension is the LLM-RL equivalent of what KnapFormer does for DiT — context-length variability in RL rollouts is the analog of sequence-length variability in video-DiT pretraining.