StreamBP: Geheugenefficiënte Exacte Backpropagatie voor Lange Sequentie Training van LLM's

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

June 3, 2025

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

cs.AI

Samenvatting

Het trainen van taalmodelen op lange sequentiedata is een veeleisende vereiste voor het verbeteren van de capaciteit van het model op complexe taken, zoals lange-ketenredenering. Naarmate de sequentielengte echter toeneemt, worden de geheugenkosten voor het opslaan van activatiewaarden enorm tijdens het Backpropagation (BP)-proces, zelfs met de toepassing van de gradient checkpointing-techniek. Om deze uitdaging aan te pakken, stellen we een geheugenefficiënte en exacte BP-methode voor, genaamd StreamBP, die een lineaire decompositie van de kettingregel uitvoert langs de sequentiedimensie op een laaggewijze manier, waardoor de geheugenkosten van activatiewaarden en logits aanzienlijk worden verminderd. De voorgestelde methode is toepasbaar op veelvoorkomende doelen zoals SFT, GRPO en DPO. Vanuit een implementatieperspectief bereikt StreamBP minder rekenkundige FLOPs en een snellere BP-snelheid door gebruik te maken van de causale structuur van het taalmodel. In vergelijking met gradient checkpointing schaalt StreamBP de maximale sequentielengte van BP op tot 2,8-5,5 keer groter, terwijl vergelijkbare of zelfs minder BP-tijd wordt gebruikt. Merk op dat de sequentielengteschaalbaarheid van StreamBP direct kan worden overgedragen naar batchgrootte-schaalbaarheid voor het versnellen van de training. We ontwikkelen verder een communicatie-efficiënte gedistribueerde StreamBP om multi-GPU-training effectief te ondersteunen en de toepasbaarheid ervan te vergroten. Onze code kan eenvoudig worden geïntegreerd in de trainingspipeline van elk transformermodel en is beschikbaar op 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.