š¤ AI Summary
Existing context parallelism approaches struggle to efficiently handle varying training sequence lengths, resulting in low computational efficiency, high communication overhead, and load imbalance. This work proposes FCPāa flexible context parallelism paradigm that overcomes the limitations of conventional ring-based communication through block-level sequence partitioning, dynamic bin-packing for hybrid sequence scheduling, arbitrary-topology point-to-point communication, and optimized distributed attention mechanisms. These innovations enable unified and efficient scheduling of both short and long sequences while achieving balanced workload distribution. Evaluated on 256 NVIDIA GPUs, FCP demonstrates near-linear scalability and achieves a 1.13ā2.21Ć improvement in Model FLOPs Utilization (MFU) for the attention module.
š 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.