StreamBP: Memory-Efficient Exact Backpropagation for Long Sequence Training of LLMs

TL;DR

StreamBP introduces a memory-efficient, exact backpropagation method that significantly extends the maximum sequence length for training large language models, reducing memory costs and improving training speed.

Contribution

The paper presents StreamBP, a novel layer-wise, linear decomposition approach for backpropagation that reduces memory usage and computational costs for long sequence training in language models.

Findings

Scales sequence length by 2.8-5.5 times compared to gradient checkpointing.

Achieves faster backpropagation with comparable or less computational cost.

Supports multi-GPU training with a communication-efficient distributed implementation.

Abstract

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. From an implementation perspective, StreamBP achieves less computational FLOPs and faster BP speed by leveraging the causal structure of the language model. Compared to gradient checkpointing, StreamBP scales up the maximum sequence length of BP by 2.8-5.5 times larger, while using comparable or even less BP time. Note that StreamBP's sequence length scaling ability can be directly transferred to batch size scaling for accelerating training. We further develop a communication-efficient distributed StreamBP to effectively support multi-GPU training and broaden its applicability. Our code can be easily integrated into the training pipeline of any transformer models and is available at https://github.com/Ledzy/StreamBP.