本文的第一作者羅琪竣、第二作者李夢(mèng)琦為香港中文大學(xué)（深圳）計(jì)算機(jī)科學(xué)博士生，本文在上海交通大學(xué)趙磊老師、香港中文大學(xué)（深圳）李肖老師的指導(dǎo)下完成。

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

來自港中文（深圳）和上海交通大學(xué)的團(tuán)隊(duì)提出 StreamBP 算法。通過對(duì)鏈?zhǔn)椒▌t進(jìn)行線性分解和分步計(jì)算，StreamBP 將大語言模型訓(xùn)練所需的激活值內(nèi)存（logits 和 layer activation）降低至梯度檢查點(diǎn)（gradient checkpointing）的 20% 左右。

• 論文標(biāo)題：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 最大序列長(zhǎng)度為梯度檢查點(diǎn)的 2.8-5.5 倍。在相同序列長(zhǎng)度下，StreamBP 的速度和梯度檢查點(diǎn)接近甚至更快。StreamBP 適用于 SFT、GRPO、PPO 和 DPO 等常見 LLM 目標(biāo)函數(shù)。代碼已開源，可集成至現(xiàn)有訓(xùn)練代碼。

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

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

對(duì)于模型中的任意函數(shù)變換  的梯度由以下鏈?zhǔn)椒▌t計(jì)算：

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

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

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

StreamBP 的核心思想

不同于梯度檢查點(diǎn)，StreamBP 避免儲(chǔ)存單層的完整激活值，而將單層的 BP 過程進(jìn)行線性分解，序列化計(jì)算并累加。注意到對(duì)于函數(shù)變換，鏈?zhǔn)椒▌t存在以下線性分解：

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

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

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

LLM 訓(xùn)練中的 StreamBP

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

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

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

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

圖：StreamBP for SFT

圖：StreamBP for GRPO

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

圖：StreamBP for DPO

實(shí)驗(yàn)結(jié)果

我們?cè)趩螐?A800-80GB GPU 上測(cè)試了不同大小的模型，StreamBP 的最大 BP 序列長(zhǎng)度為標(biāo)準(zhǔn) BP 的 23-36 倍，梯度檢查點(diǎn)的 2.5-5.5 倍。

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

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

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

表：Qwen 3-4B 單個(gè)樣本 BP 時(shí)間，序列長(zhǎng)度為 9000。

在 Deepspeed ZeRO 分布式訓(xùn)練模式下，Distributed StreamBP 比梯度檢查點(diǎn)的最大可訓(xùn)練序列長(zhǎng)度提升了5—5.6倍。