解决大模型长文本训练中的通信效率与负载不均问题。
FlashCP: Load-Balanced Communication-Efficient Context Parallelism for LLM Training
- 采用智能分片策略减少冗余通信，动态平衡计算负载。
- 在多个数据集上实现最高1.63倍的训练加速。
- 适合需要高效训练超长文本的大模型研究者使用。
上下文并行（CP）是训练大规模长文本语言模型的关键技术，通过分割序列降低内存开销。然而，现有方法存在工作负载不均、内核效率低及因静态序列分片和键值（KV）张量通信导致的冗余通信问题。本文提出FlashCP，一种负载均衡且通信高效的CP训练框架。FlashCP引入感知分片的通信机制，消除冗余的KV通信，并提出创新的Whole-Doc分片策略，在保持负载均衡的同时最大化通信节省。为有效融合Whole-Doc与Per-Doc分片，FlashCP进一步设计启发式算法以搜索近似最优分片方案。大量实验表明，FlashCP在多种数据集上相较当前最优CP框架最高提升1.63倍训练速度。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。