StreamBP: Memory-Efficient Exact Backpropagation for Long Sequence
  Training of LLMs
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.
683fc07a1de14546d5decf19
Qijun Luo
683fc07a1de14546d5decf1a
Mengqi Li
claimed_verified
2025-06-04T08:55:10.607Z
65a521af90b5e87bcd343828
/avatars/3bbe83c1ba47df17d9c05a049147e5cc.svg
Mengqi Li
user
Kullpar
683fc07a1de14546d5decf1b
Lei Zhao
683fc07a1de14546d5decf1c
Xiao Li
StreamBP: Memory-Efficient Exact Backpropagation for Long Sequence
  Training of LLMs
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.
StreamBP, a memory-efficient and exact backpropagation method, decomposes the chain rule to reduce memory costs, enabling longer sequence lengths and faster training speeds for language models compared to gradient checkpointing.
backpropagation (BP) memory-efficient exact BP gradient checkpointing chain rule sequence dimension layer-wise activation values logits SFT GRPO DPO computational FLOPs BP speed causal structure language model multi-GPU training distributed StreamBP