Equal token counts are not equal workloads in long-context training. A packed sequence can contain one long document or many short documents, and an attention mask prevents tokens from attending across document boundaries. Tokens near the end of a long document therefore perform more attention work than tokens at the same absolute position in a collection of short documents. Fixed sequence length equalizes memory shape, but not arithmetic intensity.
Meta’s trace from an internal 405B training job on 8,000 H100 GPUs makes the cost visible. With a 128K context window, the slowest observed GPU took 1.44 times as long as its peers for attention computation[1]. Synchronous collectives and pipeline dependencies propagate that delay. The useful capacity of a large training system is consequently set by the heaviest document composition, not merely by total tokens or nominal FLOPs.
WLB-LLM changes two points in the 4D training flow. At pipeline parallelism (PP), it lets micro-batches have different token lengths so their predicted total latency can be similar. At context parallelism (CP), it can shard each packed document separately, then selects that finer layout only when its balance benefit exceeds the attention kernel’s efficiency loss. The design treats input composition as a runtime scheduling variable rather than immutable loader output.
Two imbalance sources survive equal-token partitioning
The framework combines data, pipeline, context, and tensor parallelism. Data-parallel workers receive global batches. Pipeline stages execute a sequence of micro-batches. Context parallelism divides long sequences among GPUs, while tensor parallelism partitions model operations. Equal tokens appear at every boundary, yet the synchronization rule differs at each level.
At the PP level, a micro-batch containing one long document has more attention work than a micro-batch filled with shorter documents. All pipeline stages process the same micro-batch composition, and the slowest micro-batch traverses the entire pipeline dependency chain. Its delay is added to the forward and backward work that remains on the first stage. The pipeline can therefore amplify a skew that began in packing.
At the CP level, the conventional layout cuts the whole sequence into twice as many chunks as CP workers, then assigns symmetric chunk pairs. This balances one uninterrupted document because early and late attention regions are paired. Packing several documents changes where attention restarts. A chunk with several document tails can carry more query-key work than another chunk with the same token count. Tensor-parallel workers inside a CP worker still see the same gathered sequence chunk, so the imbalance is not repaired at the TP level.
Variable length converts spare memory into balancing room
Fixed-length repacking has a hard limit. If one document already fills the context window, no collection of shorter documents can match its quadratic attention work without exceeding that token limit. Expanding the packing window across several global batches improves the mathematical balance but reorders more training examples. In the paper’s 550M pretraining experiment, larger windows increased final training loss, establishing that scheduler freedom has a model-quality cost.
WLB-LLM instead permits different micro-batch lengths up to a memory-derived maximum. Its objective predicts attention latency from document lengths and adds the latency of GEMMs, elementwise work, and collectives. A collection of short documents can exceed the nominal context length when memory allows, because its masked attention work remains lower than that of one long document. The extra non-attention work raises total latency toward the outlier without forcing unrelated documents across a large reordering window.
The runtime algorithm sorts pending documents by length and greedily assigns each to the micro-batch with the least predicted work. If that placement would exceed the memory bound, it tries the shortest current micro-batch; documents that still do not fit remain for the next iteration. Offline profiles supply the attention and other-operation latency functions. This is a configuration-sensitive model, not a universal token formula. A new GPU, kernel, precision, or communication mapping requires refreshed profiles.
Rare outliers are delayed, not globally shuffled
One global batch may not contain enough comparable long documents to balance every micro-batch. WLB-LLM places unusually long documents into length-banded FIFO queues. A queue releases documents only after it has enough entries to give every micro-batch a similar outlier. The thresholds trade balance against waiting time: narrow bands match work more closely but accumulate entries more slowly.
The authors tune thresholds on a sample of the training corpus, optimizing imbalance subject to a token-delay constraint. With two outlier queues, the reported imbalance metric reached 1.05 and packing took 20 ms per batch, less than 0.65% of step latency. A fixed-length integer-programming solver reached comparable balance only with much higher runtime, exceeding 25 seconds per batch when solving across four global batches.
Average token delay was 0.5 training iteration. The paper reports that WLB-LLM’s loss curve followed fixed-length packing within one global batch, whereas repacking across more batches increased loss. This evidence supports the evaluated pretraining setup, but it does not prove that every optimizer, curriculum, rare-domain mixture, or exact-resumption policy tolerates the same delay. A deployment should track per-domain delay distributions, not only the global average.
Per-document sharding balances work but can waste tiles
At the CP boundary, WLB-LLM divides every document into twice the CP group size and gives each worker symmetric pieces from each document. It distributes remainders round-robin, avoiding explicit padding while keeping equal token totals. Because each worker receives comparable early and late regions from every document, both token count and attention work become balanced.
Finer sharding can nevertheless slow the kernel. FlashAttention-style implementations operate on tiles, so a short query fragment can occupy the same tile as a longer fragment and waste computation. On Hopper GPUs, sufficiently large query regions can also reuse key-value loads through Tensor Memory Accelerator multicast and L2. Breaking one document into small pieces reduces that reuse and lowers achieved TFLOPs. Perfect work balance is not the same as minimum step time.
WLB-LLM therefore evaluates two candidates for every micro-batch. It derives query and key-value shapes for per-sequence and per-document sharding, rounds them to the kernel tile, estimates FLOPs, and divides by achieved TFLOPs from offline profiling. The lower predicted maximum CP-worker latency wins. In the 7B-128K breakdown, per-document sharding alone delivered 1.02 times, while adaptive selection raised the CP contribution to 1.05 times. The PP packing and delay mechanism accounted for the larger 1.28-times component; combining both reached 1.33 times.
The reported gain grows with context, not model size
The evaluation used 32 nodes, each with eight H100 SXM 80GB GPUs connected by NVLink, with RoCE across nodes. Internal LLaMA-like models ranged from 550M to 70B parameters. Context windows were 64K and 128K, and each scale used a specific tensor, context, pipeline, and data parallel configuration. The largest tested case used 256 GPUs, not the 8,000-GPU trace that motivated the design.
Against the internal Plain-4D framework, WLB-LLM averaged 1.23 times and exceeded the Fixed-4D comparison by 1.19 times. Fixed-4D improved only 1.03 times on average because it repacked within one global batch and used one CP sharding choice for an entire run. For the 7B model, speedup increased from 1.03 at 32K context to 1.40 at 160K. Across the main configurations, moving from 64K to 128K raised the average from 1.15 to 1.30 times.
Larger models showed smaller relative gains because communication occupied more of step time while the optimization targets attention-work imbalance. Longer contexts increase both attention’s share and the probability of an outlier document, so they enlarge the addressable fraction. This distinction matters for capacity planning. The average 1.23-times result should not be applied to short-context, communication-bound, or differently packed jobs without measuring their imbalance profile.
Scheduling metadata becomes part of the training contract
WLB-LLM shows that distributed-training balance needs document boundaries, not just tensor shapes. A practical rollout must preserve those boundaries through the loader, packing metadata, CP partitioner, and attention kernel. It also needs deterministic queue state in checkpoints so restart does not silently change example order. Offline latency profiles must be versioned with GPU architecture, kernel build, precision, and parallel mapping.
Mixture-of-Experts adds another imbalance source after tokens are routed to experts. The paper argues that its packing and sharding are compatible with dropless expert routing because they do not change gating decisions. That compatibility does not remove the need to measure expert hot spots. Input balance and expert balance are separate schedulers whose worst cases may coincide.
The procurement lesson is narrower than “buy 23% fewer GPUs.” WLB-LLM improves completed training work per installed GPU when long-document composition creates synchronized idle time. Operators should first measure the slowest-to-average attention ratio at PP and CP boundaries, the fraction of step time in attention, token-delay distribution, packing overhead, and model-quality trajectory. If those counters reproduce the paper’s failure mode, document-aware scheduling can recover capacity without changing model mathematics. If communication or expert routing dominates, another bottleneck sets the return.
Source and copyright notice
This article is an editorial analysis by Silicon & Systems. It restates the architecture, evaluation, and limitations in our own words. No source sentence, table, or figure is reproduced; the figure was created for this article. The paper is available from the USENIX OSDI 2025 presentation page. Copyright remains with the authors, 2025.