FlashCP: Load-Balanced Communication-Efficient Context Parallelism for LLM Training
This work addresses workload imbalance and communication inefficiency in context parallelism for training long-context LLMs, offering a practical speedup for large-scale training.
FlashCP introduces a load-balanced and communication-efficient framework for context parallelism in LLM training, achieving up to 1.63x speedup over existing methods by eliminating redundant KV communication and optimizing sharding strategies.
Context parallelism (CP) is essential for training large-scale, long-context language models, as it partitions sequences to reduce memory overhead. However, existing CP methods suffer from workload imbalance, inefficient kernels, and redundant communication due to static sequence sharding and key-value (KV) tensor communication. We present FlashCP, a load-balanced and communication-efficient framework for CP training. FlashCP introduces a sharding-aware communication mechanism to eliminate redundant KV communication and proposes a novel Whole-Doc sharding strategy that maximizes communication savings while maintaining balanced workloads. To efficiently combine Whole-Doc and Per-Doc sharding, FlashCP further designs a heuristic algorithm to search for near-optimal sharding plans. Extensive experiments show that FlashCP achieves up to 1.63x speedup over state-of-the-art CP frameworks across diverse datasets.