Increasing the context length of large language models (LLMs)  unlocks fundamentally new capabilities, but also significantly increases the memory footprints of training.
Previous model-parallel systems such as Megatron-LM partition and compute different attention heads in parallel, resulting in large communication volumes,
so they cannot scale beyond the number of attention heads, thereby hindering its adoption.
In this paper, we introduce a new approach,LightSeq, for long-context LLMs training.LightSeqhas many notable advantages. First,LightSeqpartitions over the sequence dimension, hence is agnostic to model architectures and readily applicable for models with varying numbers of attention heads, such as Multi-Head, Multi-Query and Grouped-Query attention. Second,LightSeqnot only requires up to 4.7less communication than Megatron-LM on popular LLMs but also overlaps the communication with computation. To further reduce the training time,LightSeqfeatures a novel gradient checkpointing scheme to bypass an forward computation for memory-efficient attention.
We evaluateLightSeqon Llama-7B and its variants with sequence lengths from 32K to 512K. Through comprehensive experiments on single and cross-node training, we show thatLightSeqachieves up to 1.24-2.01end-to-end speedup, and a 2-8longer sequence length on models with fewer heads, compared to Megatron-LM. Codes will be available athttps://github.com/RulinShao/LightSeq.