Unleashing Scalable Context Parallelism for Foundation Models Pre-Training via FCP

Abstract: Context parallelism (CP) has been widely adopted to support the growing context length in foundation model pretraining. However, existing designs fail to handle the large variation in sequence length from training datasets, resulting in suboptimal performance. These methods often over-shard short sequences, leading to compute inefficiency and excessive communication, or process long and short sequences separately without proper bin-packing, causing workload imbalance. In this paper, we propose FCP, a flexible context parallelism paradigm that shards and schedules sequences at block-level granularity. Instead of relying on rigid communication topologies such as ring, FCP enables arbitrary peer-to-peer communication, allowing flexible placement of sequence blocks across workers. By bin-packing blocks from both short and long sequences, FCP achieves both high compute efficiency and balanced workload distribution. Extensive evaluations show that FCP attains near-linear scalability on up to 256 NVIDIA GPUs, with 1.13x-2.21x improvement in the attention MFU.

• The paper introduces FCP, a flexible context-parallelism system that combines 4K-token block sharding, LPT-based workload balancing, and matching-based peer-to-peer scheduling to replace restrictive ring topologies.
• FCP keeps workload imbalance below 5% and improves attention MFU by 1.13×–2.21% over competing methods, achieving near-linear weak scaling across up to 256 GPUs on heterogeneous long-context workloads.
• The system hides layout reshuffling and communication behind computation, but its benefits depend on high-bandwidth network fabrics and remain to be tested for irregular sparse attention patterns and adaptive block sizes.

Motivation and problem statement

Context parallelism (CP) shards attention computation across GPUs to support the long sequences that now dominate foundation-model pre-training, where a single 1080p image yields thousands of patch tokens and a minute of 30 fps video can expand to millions of tokens. The central obstacle addressed in this paper is that real pre-training corpora exhibit highly diverse, long-tailed sequence-length distributions — the authors' internal traces with a maximum length of 512K approximately follow a lognormal distribution, mixing short text samples with extremely long multimodal ones. Under such heterogeneity, existing CP designs are systematically suboptimal for one of two reasons: they either over-shard short sequences into blocks too small to saturate Tensor Cores while still paying communication costs, or they isolate short- and long-sequence groups on separate workers without resource sharing, so that quadratic attention compute collides with linearly allocated GPUs.

The paper formalizes CP scheduling as a choice of a sharding function $G$ mapping each sequence to blocks and an assignment function $M$ mapping blocks to workers, with end-to-end time $T = \max_i (\eta_i \cdot Comp(w_i))$, where $f(\cdot)$ captures both FLOPs and block-size-dependent compute efficiency and $\eta_i \geq 1$ is the network overlap factor. The authors state that this optimization is NP-complete, which motivates all prior rule-based simplifications. Their key structural observation is that every existing design retains a ring-based assignment policy — either monolithic or as heterogeneous sub-rings — and that this "gilded cage" of symmetric ring topology, not the sharding heuristic alone, is what fundamentally restricts the search space. Because per-block computation scales quadratically with block size while communication scales linearly, sufficiently large blocks make arbitrary peer-to-peer placement affordable: the bandwidth required for communication to match computation falls as blocks grow. This is the opportunity FCP exploits.

Design

FCP replaces fixed rings with fine-grained, block-wise scheduling built on three components.

Block distributor. Each sequence is partitioned into fixed-size blocks (4K tokens by default) regardless of its original length; sequences shorter than the block size are packed via the varlen kernel API. Assignment uses a variant of Longest Processing Time (LPT) scheduling formulated as multi-dimensional bin-packing: blocks are sorted by normalized compute/memory cost and greedily placed on the least-loaded worker under a per-worker memory constraint, in $O(K \log N)$ time. Zig-Zag ordering within each sequence reduces communication balance to computation balance under causal masks.

Communication planner. Arbitrary peer-to-peer placement creates irregular traffic that must be overlapped with computation ($\eta = 1$). FCP decomposes execution into block-level sub-stages of pulling remote KV blocks, computing attention, and pushing local blocks. To order these sub-stages without congestion, it models data flow as an undirected bipartite graph over $N$ send and $N$ receive nodes, and proves two lemmas: a single congestion-free stage is exactly a matching on this graph, and a bipartite graph of maximal degree $\Delta$ requires at least $M$0 disjoint matchings. Via Hall's theorem, the edge set decomposes into exactly $M$1 matchings, computed by Hopcroft–Karp in $M$2 time once per batch — seconds at the scale of hundreds of workers, parallelizable across CPUs. A bottom-up coalescer then merges consecutive sub-stages (degree 16 by default), decoupling scheduling granularity from execution granularity and recovering kernel efficiency without introducing hotspots.

Transparent reshuffler. Rather than requiring intrusive changes to dataloaders and positional embeddings, FCP reshuffles the user-provided layout into the workload-aware layout on entry to each layer's attention module and restores it afterward. Because reshuffling volume is bounded by total context length while computation grows quadratically, the all-to-all traffic is hidden behind local attention computation scheduled at pipeline boundaries.

Implementation-wise, FCP comprises roughly 4K lines of Python with minor FlashAttention3 modifications, uses CUDA Green Contexts to partition SMs between compute and communication (6–8 SMs suffice since network, not HBM, is the bottleneck), NCCL group peer-to-peer primitives topology-aware of rail-optimized fabrics, and a three-buffer pipeline.

Evaluation

Experiments use the Llama-3-70B configuration (64 QO heads, 8 KV heads, head dim 128) on two anonymized GPU types across up to 256 GPUs, with 32K tokens per GPU and real trace-derived workloads plus synthetic lognormal and bimodal distributions. Baselines are Ring Attention, ByteScale, WLB-LLM (with an oracle estimator), and MagiAttention.

The headline results are strong. FCP holds workload imbalance below 5% at all scales, versus up to 17% communication imbalance for MagiAttention and up to 70% compute imbalance for ByteScale on long-tailed inputs. With perfect load balance assumed, FCP sustains above 90% attention MFU, against severe degradation for Ring Attention from over-sharding short sequences. End-to-end, FCP outperforms all baselines by 1.13×–2.21× in module-level attention MFU with near-linear weak scaling to 256 GPUs. Ablations attribute gains progressively: starting from a 0.29 forward MFU baseline, adding block-level pipelining, the congestion-free solver, the coalescer, and the reshuffler raises forward MFU to 0.75 (and backward from 0.37 to 0.74) on 128 GPU-X. Sensitivity tests identify 4K as the best block size — about 7% better than 2K or 6K — and confirm robustness across per-GPU token counts and across GPU-Y with FlashAttention-4, where FCP exceeds 70% of single-GPU FA4 MFU.

Limitations and open questions

The paper is explicit about several constraints. First, FCP's performance depends on network topology: its coalesced peer-to-peer traffic effectively resembles all-to-all, so it generalizes well on fat-tree or rail-optimized InfiniBand/RoCE fabrics but performs poorly on torus-based topologies such as TPU v3 clusters. Second, block size selection involves a three-way trade-off among kernel saturation, overlap feasibility, and balance granularity; the authors argue hardware-driven defaults suffice for production workloads but concede only marginal benefit from further tuning, leaving open whether adaptive per-batch block sizing could do better. Third, the design and evaluation cover causal and non-causal masks only; irregular and block-sparse mask patterns are deferred, though the authors note the block abstraction could extend to sparsity maps. Finally, the evaluation anonymizes exact GPU models and cluster scales for confidentiality, and the strongest comparisons rely on reimplementations of WLB-LLM and ByteScale rather than their official end-to-end systems.

Conclusion

FCP demonstrates that removing the ring-topology constraint from context parallelism — replacing it with fixed-size block sharding, LPT-based bin-packing, and provably congestion-free matching-based communication ordering — yields both high single-GPU compute efficiency and near-perfect cluster-level load balance under long-tailed sequence-length distributions. Its modular integration with FSDP, TP, EP, and SP, together with demonstrated scalability to 256 GPUs and consistent 1.13×–2.21× MFU improvements over state-of-the-art baselines, establishes flexible peer-to-peer CP as a practical alternative to ring-based designs for heterogeneous long-context pre-training.