StreamBP: 長シーケンスLLMトレーニングのためのメモリ効率の良い正確な誤差逆伝播法
StreamBP: Memory-Efficient Exact Backpropagation for Long Sequence Training of LLMs
June 3, 2025
著者: Qijun Luo, Mengqi Li, Lei Zhao, Xiao Li
cs.AI
要旨
長いシーケンスデータを用いた言語モデルのトレーニングは、複雑なタスク（例：長い連鎖推論）におけるモデルの能力を向上させる上で重要な要件です。しかし、シーケンス長が増大するにつれて、バックプロパゲーション（BP）プロセスにおける活性化値のメモリコストが膨大になり、勾配チェックポイント技術を適用した場合でもその問題は残ります。この課題に対処するため、我々はメモリ効率が高く正確なBP手法であるStreamBPを提案します。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.