StreamBP: Memory-Efficient Exact Backpropagation for Long Sequence Training of LLMs
Qijun Luo, Mengqi Li, Lei Zhao, Xiao Li
Training language models on long sequence data is a demanding requirement for enhancing the model's capability on complex tasks, e.g., long-chain reasoning. However, as the sequence length scales up, the memory cost for storing activation values becomes huge during the Backpropagation (BP) process, even with the application of gradient checkpointing technique. To tackle this challenge, we propose a *memory-efficient* and *exact* BP method called **StreamBP**, which performs a linear decomposition of the chain rule along the sequence dimension in a layer-wise manner, significantly reducing the memory cost of activation values and logits. The proposed method is applicable to common objectives such as SFT, GRPO, and DPO. Moreover, by leveraging the causal structure of LLMs, StreamBP reduces the computational FLOPs and allows faster BP speed compared to the gradient checkpointing baseline. StreamBP scales up the maximum sequence length of BP by $2.8-5.5 \times$ larger, and reduces BP time cost by approximately $30$% compared to gradient checkpointing. We further develop a communication-efficient distributed StreamBP to effectively support multi-GPU training and broaden its applicability.