本文的第一作者羅琪竣、第二作者李夢琦為香港中文大學（深圳）計算機科學博士生，本文在上海交通大學趙磊老師、香港中文大學（深圳）李肖老師的指導下完成。

長序列訓練對于模型的長序列推理等能力至關重要。隨著序列長度增加，訓練所需儲存的激活值快速增加，占據(jù)訓練的大部分內(nèi)存。即便使用梯度檢查點（gradient checkpointing）方法，激活值依然占據(jù)大量內(nèi)存，限制訓練所能使用的序列長度。

來自港中文（深圳）和上海交通大學的團隊提出 StreamBP 算法。通過對鏈式法則進行線性分解和分步計算，StreamBP 將大語言模型訓練所需的激活值內(nèi)存（logits 和 layer activation）降低至梯度檢查點（gradient checkpointing）的 20% 左右。

• 論文標題：StreamBP: Memory-Efficient Exact Backpropagation for Long Sequence Training of LLMs
• 論文：https://arxiv.org/abs/2506.03077
  
• 代碼：https://github.com/Ledzy/StreamBP

在相同內(nèi)存限制下，StreamBP 最大序列長度為梯度檢查點的 2.8-5.5 倍。在相同序列長度下，StreamBP 的速度和梯度檢查點接近甚至更快。StreamBP 適用于 SFT、GRPO、PPO 和 DPO 等常見 LLM 目標函數(shù)。代碼已開源，可集成至現(xiàn)有訓練代碼。

激活值內(nèi)存和梯度檢查點

在反向傳播（Backpropagation, BP）的過程中，計算模型梯度需要用到模型的中間輸出（激活值）。舉例來說，對于模型中的線性變換的梯度為，因而計算的梯度時需要儲存相應的激活值。

對于模型中的任意函數(shù)變換  的梯度由以下鏈式法則計算：

其中 L 為目標函數(shù)，為 Jacobian 矩陣。為了計算以上 Jacobian-vector product，需要在模型 forward 時儲存函數(shù)變換的中間值（激活值），其內(nèi)存消耗與 batch size、序列長度以及中間值維度正相關。

為了減少激活值的內(nèi)存消耗，梯度檢查點（gradient checkpointing）方法在 forward 時只儲存每一層網(wǎng)絡的輸入，而不儲存該層的中間值。在 backward 至該層時，將重新 forward 此層輸入來計算得到該層激活值。使用梯度檢查點時儲存的激活值包括：

• 所有層的輸入，一般為激活值內(nèi)存的 5%-15%。
• 單層的完整激活值，占據(jù)超過 85% 的激活值內(nèi)存。

StreamBP 的核心思想

不同于梯度檢查點，StreamBP 避免儲存單層的完整激活值，而將單層的 BP 過程進行線性分解，序列化計算并累加。注意到對于函數(shù)變換，鏈式法則存在以下線性分解：

StreamBP 基于以下觀察：對于 LLM 中的大部分函數(shù)變換，如 Transformer 層、lmhead 層，可通過策略性地將輸出分塊，使得計算塊 Jacobian-vector product 所需的激活值遠小于計算完整的 Jacobian-vector product?；谠撚^察，StreamBP 依次計算上式中 D 個塊的 Jacobian-vector product 并累加，得到準確的梯度。

為了計算塊 Jacobian-vector product，需要分析輸入和輸出的相關性，每次 forward 塊輸入 得到塊輸出，建立對應子計算圖。以簡單的線性變換 為例，輸出和輸入在行維度上一一對應。StreamBP 按行分塊，每次計算單行的 Jacobian-vector product 并累加。下圖對比了標準 BP 和 StreamBP 在上述線性變換下的實現(xiàn)：

D 步累加得到的和即為和準確梯度。相比于標準 BP，StreamBP 僅需儲存和，且總計算 FLOPs 相同。下表為 StreamBP 和標準 BP 的內(nèi)存和時間對比：

LLM 訓練中的 StreamBP

StreamBP 應用于 LLM 中的 Transformer 層和 lmhead 層，分別用于降低層激活值和 logits 的內(nèi)存消耗。

與線性變換不同，由于 Transformer 層存在注意力機制，塊輸出并非僅由對應位置的塊輸入決定，而與該塊及以前所有位置的輸入都有關。StreamBP 利用只與塊有關的性質(zhì)，建立了如下計算圖：

StreamBP 所需儲存的激活值和注意力掩碼（橙色）大幅低于梯度檢查點（橙色 + 白色部分）。

對于 lmhead 層，當以 SFT 或 GRPO 為目標函數(shù)時，觀察到不同位置的 logits 對于目標函數(shù)的影響相互獨立。因此，StreamBP 從序列維度分塊，每次計算單塊損失函數(shù)的梯度，從而只需儲存單塊 logits 和 logits 梯度。

圖：StreamBP for SFT

圖：StreamBP for GRPO

對于 DPO，由于非線性 sigmoid 函數(shù)的存在，每個位置的 logits 對于目標函數(shù)的影響并不獨立。StreamBP 利用 logits 梯度在序列維度的獨立性，分塊進行梯度計算。

圖：StreamBP for DPO

實驗結果

我們在單張 A800-80GB GPU 上測試了不同大小的模型，StreamBP 的最大 BP 序列長度為標準 BP 的 23-36 倍，梯度檢查點的 2.5-5.5 倍。

圖：不同序列長度下的 BP 峰值內(nèi)存

在現(xiàn)有 Transformers 框架下，StreamBP 的實現(xiàn)可避免計算掩碼部分的 pre-attention score（見論文 3.2.2 部分），在長序列訓練下相較于梯度檢查點實現(xiàn)了加速。

通過使用 StreamBP，不同目標函數(shù)下最大的序列長度得到了大幅提升。在同樣的序列長度下，StreamBP 允許更大的批處理大小以加速訓練。

表：Qwen 3-4B 單個樣本 BP 時間，序列長度為 9000。

在 Deepspeed ZeRO 分布式訓練模式下，Distributed StreamBP 比梯度檢查點的最大可訓練序列長度提升了5—5.6倍。