StreamBP: Retropropagação Exata com Eficiência de Memória para Treinamento de Longas Sequências em LLMs

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

June 3, 2025

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

cs.AI

Resumo

Treinar modelos de linguagem em dados de sequência longa é uma exigência desafiadora para aprimorar a capacidade do modelo em tarefas complexas, como raciocínio de cadeia longa. No entanto, à medida que o comprimento da sequência aumenta, o custo de memória para armazenar valores de ativação torna-se enorme durante o processo de Retropropagação (BP), mesmo com a aplicação da técnica de checkpoint de gradiente. Para enfrentar esse desafio, propomos um método de BP eficiente em memória e exato chamado StreamBP, que realiza uma decomposição linear da regra da cadeia ao longo da dimensão da sequência de maneira camada por camada, reduzindo significativamente o custo de memória dos valores de ativação e logits. O método proposto é aplicável a objetivos comuns, como SFT, GRPO e DPO. Do ponto de vista de implementação, o StreamBP alcança menos FLOPs computacionais e uma velocidade de BP mais rápida ao aproveitar a estrutura causal do modelo de linguagem. Em comparação com o checkpoint de gradiente, o StreamBP aumenta o comprimento máximo da sequência de BP em 2,8 a 5,5 vezes, enquanto utiliza um tempo de BP comparável ou até menor. Observe que a capacidade de escalonamento de comprimento de sequência do StreamBP pode ser diretamente transferida para o escalonamento do tamanho do lote para acelerar o treinamento. Desenvolvemos ainda um StreamBP distribuído com comunicação eficiente para apoiar efetivamente o treinamento em multi-GPU e ampliar sua aplicabilidade. Nosso código pode ser facilmente integrado ao pipeline de treinamento de qualquer modelo transformer e está disponível em 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.