StreamBP: Эффективное по памяти точное обратное распространение для обучения языковых моделей на длинных последовательностях
StreamBP: Memory-Efficient Exact Backpropagation for Long Sequence Training of LLMs
June 3, 2025
Авторы: Qijun Luo, Mengqi Li, Lei Zhao, Xiao Li
cs.AI
Аннотация
Обучение языковых моделей на данных с длинными последовательностями является важным требованием для повышения их способности выполнять сложные задачи, такие как длинные цепочки рассуждений. Однако с увеличением длины последовательности затраты памяти на хранение значений активаций становятся огромными в процессе обратного распространения ошибки (Backpropagation, BP), даже при использовании техники градиентного чекпоинтинга. Для решения этой проблемы мы предлагаем эффективный по памяти и точный метод BP под названием StreamBP, который выполняет линейное разложение цепного правила вдоль последовательности послойно, значительно снижая затраты памяти на значения активаций и логиты. Предложенный метод применим к распространенным задачам, таким как SFT, GRPO и DPO. С точки зрения реализации, StreamBP достигает меньшего количества вычислительных операций (FLOPs) и более высокой скорости BP за счет использования причинной структуры языковой модели. По сравнению с градиентным чекпоинтингом, StreamBP увеличивает максимальную длину последовательности для BP в 2,8–5,5 раз, при этом используя сопоставимое или даже меньшее время BP. Отметим, что способность StreamBP масштабировать длину последовательности может быть напрямую перенесена на масштабирование размера батча для ускорения обучения. Мы также разработали коммуникационно-эффективную распределенную версию StreamBP, которая эффективно поддерживает обучение на нескольких GPU и расширяет область его применения. Наш код легко интегрируется в конвейер обучения любых трансформерных моделей и доступен по адресу 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.