StreamBP: Backpropagazione Esatta Efficiente in Memoria per l'Addestramento su Sequenze Lunghe di LLM
StreamBP: Memory-Efficient Exact Backpropagation for Long Sequence Training of LLMs
June 3, 2025
Autori: Qijun Luo, Mengqi Li, Lei Zhao, Xiao Li
cs.AI
Abstract
L'addestramento di modelli linguistici su dati a sequenza lunga è un requisito impegnativo per migliorare le capacità del modello su compiti complessi, ad esempio il ragionamento a catena lunga. Tuttavia, man mano che la lunghezza della sequenza aumenta, il costo della memoria per memorizzare i valori di attivazione diventa enorme durante il processo di Backpropagation (BP), anche con l'applicazione della tecnica di gradient checkpointing. Per affrontare questa sfida, proponiamo un metodo BP efficiente in termini di memoria e preciso chiamato StreamBP, che esegue una decomposizione lineare della regola della catena lungo la dimensione della sequenza in modo stratificato, riducendo significativamente il costo della memoria per i valori di attivazione e i logit. Il metodo proposto è applicabile a obiettivi comuni come SFT, GRPO e DPO. Dal punto di vista dell'implementazione, StreamBP ottiene meno FLOP computazionali e una velocità BP più rapida sfruttando la struttura causale del modello linguistico. Rispetto al gradient checkpointing, StreamBP aumenta la lunghezza massima della sequenza BP di 2,8-5,5 volte, utilizzando un tempo BP comparabile o addirittura inferiore. Si noti che la capacità di scalabilità della lunghezza della sequenza di StreamBP può essere trasferita direttamente alla scalabilità della dimensione del batch per accelerare l'addestramento. Abbiamo inoltre sviluppato una versione distribuita di StreamBP efficiente dal punto di vista della comunicazione per supportare efficacemente l'addestramento multi-GPU e ampliarne l'applicabilità. Il nostro codice può essere facilmente integrato nella pipeline di addestramento di qualsiasi modello transformer ed è disponibile all'indirizzo https://github.com/Ledzy/StreamBP.
English
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.