Search papers, labs, and topics across Lattice.
This paper introduces Batch-Aware Sequence Parallelism (BASP), a novel approach that optimizes sequence parallelism in large language model (LLM) training by leveraging batch structure to minimize communication overhead. By partitioning GPUs into disjoint sequence-parallel groups based on micro-batch size, BASP effectively reduces the size of all-to-all communication groups, leading to significant improvements in training efficiency. Experimental results demonstrate that BASP enhances end-to-end training time by 1.17 to 1.31 times for Llama and Qwen models without compromising model accuracy or memory usage.
BASP slashes communication overhead in LLM training, boosting efficiency by over 30% while maintaining model performance.
Long-context reasoning for large language models (LLMs) is becoming increasingly important, but training over long sequences remains challenging due to massive memory and communication requirements. Sequence parallelism has emerged as an essential technique for addressing bottlenecks in long sequence LLM training. However, we observe that existing sequence parallelism methods are batch-agnostic and apply uniform sequence partitioning across all batch sizes, resulting in inefficient communication. In this paper, we introduce Batch- Aware Sequence Parallelism (BASP), a sequence parallelism approach that leverages batch structure to reduce communication overhead. BASP exploits batch structure by partitioning GPUs into disjoint sequence-parallel groups according to the micro- batch size. This design reduces the all-to-all communication group size, thereby localizing communication and improving training efficiency. Experimental results on an NVIDIA A100 cluster show that BASP improves end-to-end training time by up to 1.17 - 1.31x in Llama and Qwen models compared to standard sequence parallel baselines, while preserving identical model accuracy and memory usage.
