Libra: Taming Attention Workload Skew in Long-Context LLM Training with Bounded Sequence Pool

2026-07-25

AI Summary

This work addresses severe load imbalance in long-context large language model training, where the quadratic complexity of attention computation with respect to sequence length leads to straggling in data parallelism and pipeline bubbles. To mitigate this, the authors propose a bounded sequence pooling mechanism that decouples load balancing from data parallelism. Combined with variance-reduced sequence placement, SH-Tiles (chunked attention pooling), communication-computation overlap, and a pluggable sampler, the approach enables efficient scheduling without expanding the communication domain. Evaluated on Qwen3-Turbo training, the method achieves up to a 2.54× end-to-end throughput improvement and accelerates the worst-case attention straggler steps by 3.14×, demonstrating stable operation in production environments over hundreds of thousands of GPU hours.

Abstract

Long-context LLM training suffers from a load-balancing problem that sequence packing does not solve. Packing samples into fixed-token sequences balances memory and linear-cost operators, but the dominant attention cost scales with the sum of squared sequence lengths. Thus, equally sized packed sequences drawn from a long-tailed corpus can carry substantially different attention workloads, creating data-parallel stragglers and pipeline bubbles. Existing approaches either balance at the granularity of sequences or microbatches, where an outlier can dominate an assignment, or disaggregate attention over a global worker pool whose communication domain grows with the data-parallel (DP) degree. We present Libra, which operationalizes the law of large numbers (LLN) as a scaling principle for load balancing: the attention-balancing pool need not grow with the DP degree. Libra groups packed sequences and their CP groups into fixed-size sequence pools. As DP scales out, Libra adds pools rather than enlarging each one, bounding every attention exchange. Variance-Reduced Sequence Placement makes this effective for finite, long-tailed workloads by co-locating sequences with complementary attention workloads to reduce residual inter-pool skew. Within each pool, Tiled Attention Pooling dispatches sequence-head SH-Tiles across GPUs, while a pipelined runtime overlaps tile exchange with attention. Libra exposes a drop-in context-parallel attention operator and a pluggable data sampler, requiring no changes to model layers, optimizers, or pipeline schedules. On Qwen3-Turbo training with 256K- and 1M-token workloads, Libra improves end-to-end throughput by up to 2.54x over Ulysses, with up to 3.14x worst-step straggler-attention speedup in microbenchmarks. Libra has run for hundreds of thousands of GPU-hours in production on jobs spanning 32K to 1M tokens while preserving training semantics.

Problem

• long-context LLM training
• attention workload skew
• load balancing
• sequence packing
• data-parallel stragglers

Innovation

• load balancing
• long-context LLM training
• attention workload skew
• sequence packing
• context parallelism