Pipelining the Memory Hierarchy for Ultra-Long-Context Language Model Training

Long-context language models are attractive because they can read whole documents, code repositories, conversations, and scientific sequences. However, a long inference window is not the same as a model that was trained to use millions of tokens reliably. Direct long-context training is blocked by temporary memory spikes from activations, attention buffers, communication workspaces, and the final loss computation, even when model parameters and optimizer states are already sharded. Our proposed method, Fully Pipelined Distributed Transformer (FPDT), treats the long sequence problem as a memory-hierarchy scheduling problem. It chunks each sequence before distributed attention, stages inactive query/ key/value tensors in host memory and uses double buffers to overlap transfers with attention computation. FPDT democratizes long sequence training, making it accessible on mobile and commodity hardware, such as high-end local machines. In our evaluation, FPDT trains contexts up to 16x longer than prior sequence-parallel baselines on the same hardware. It reaches a 2M-token context for an 8B model on four A100-80GB GPUs, and a 4M-token context for 8B Llama 3 on eight A100-80GB GPUs while sustaining 55.7% Model FLOPs Utilization (MFU) [1].