StreamBP: Speichereffiziente exakte Backpropagation für das Training von LLMs mit langen Sequenzen

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

June 3, 2025

Autoren: Qijun Luo, Mengqi Li, Lei Zhao, Xiao Li

cs.AI

Zusammenfassung

Das Training von Sprachmodellen auf langen Sequenzdaten ist eine anspruchsvolle Anforderung, um die Fähigkeiten des Modells bei komplexen Aufgaben, wie z.B. langkettigem Schlussfolgern, zu verbessern. Allerdings steigen mit zunehmender Sequenzlänge die Speicherkosten für die Speicherung von Aktivierungswerten während des Backpropagation (BP)-Prozesses erheblich an, selbst bei Anwendung der Gradient-Checkpointing-Technik. Um diese Herausforderung zu bewältigen, schlagen wir eine speichereffiziente und exakte BP-Methode namens StreamBP vor, die eine lineare Zerlegung der Kettenregel entlang der Sequenzdimension auf Schichtebene durchführt und dadurch die Speicherkosten für Aktivierungswerte und Logits erheblich reduziert. Die vorgeschlagene Methode ist auf gängige Ziele wie SFT, GRPO und DPO anwendbar. Aus Implementierungssicht erreicht StreamBP durch die Nutzung der kausalen Struktur des Sprachmodells weniger Rechenoperationen (FLOPs) und eine schnellere BP-Geschwindigkeit. Im Vergleich zu Gradient-Checkpointing skaliert StreamBP die maximale Sequenzlänge des BP um das 2,8- bis 5,5-fache, bei vergleichbarer oder sogar geringerer BP-Zeit. Es ist zu beachten, dass die Skalierungsfähigkeit der Sequenzlänge von StreamBP direkt auf die Skalierung der Batch-Größe übertragen werden kann, um das Training zu beschleunigen. Wir entwickeln weiterhin eine kommunikationseffiziente, verteilte Version von StreamBP, um das Multi-GPU-Training effektiv zu unterstützen und die Anwendbarkeit zu erweitern. Unser Code kann leicht in den Trainingspipeline beliebiger Transformer-Modelle integriert werden und ist unter https://github.com/Ledzy/StreamBP verfügbar.

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.