StreamBP: Memory-Efficient Exact Backpropagation for Long Sequence Training of LLMs
Authors: Qijun Luo
Mengqi Li
Lei Zhao
Xiao Li
Venue: NeurIPS 2025 poster ()
Keywords: Memory efficiency
long sequence training
training reasoning LLMs
Primary area: deep_learning

Abstract: 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.

StreamBP: Memory-Efficien Exact Backpropagation for Long Sequence Training of LLMs
Qijun Luo'
Mengqi Li
Lei Zhao?
Xiao Li
'The Chinese University of Hong Kong, Shenzhen
Shanghai Jiao Tong University
{qijunluo menggilin) @link. cuhk, edu. en, 1.zhao0s ee..edt..uuue..e.... cn, lixiao@cuhk. edu. en
edu
Abstract
Training language models on long sequence data is a demanding requirement for enhancing the model's capability 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, application of gradient checkpointing technique. To tackle this challenge, we propose 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 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.5x 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 distributed StreamBP to effectively support
with the
develop a communication--ffice multi-GPU training and broaden its applicability. Our code can be easily integrated into the training pipeline of any transformer models and available at https: //gi //github. cwm/codm////eeeee Seadzy/Stream:e
Introduction
Large language models (LLMs) [1, 6, 32] have been regarded as powerful approach toward general artificial intelligence [3]. Recently, there has been growing interest in training LLMs for reasoning tasks, leading to significant improvements on math and code benchmarks as shown by OpenAl-ol and DeepSeek-R1 [12. 8]. However, such models often require extremely long input sequences due to the inclusion of detailed reasoning traces [8, 35, 24, 33, 34], which results in substantial GPU memory consumption during backpropagation (BP) when training reasoning models using either reinforcement learning (RL) or supervised finetuing (SFT). This considerable memory usage mainly stems from storing intermediate activations during the BP process. This work aims to provide an efficient solution to such memory issue.
Background and related works. Due to rapid growth of research in this field, provide non-exhaustive background overview on reasoning models and memory-efficien training methods.
*Corresponiing Author
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Training LLM on long reasoning traces. To incentivize long chain of thought (CoT) [29] reasoning capability, RL is one of most popular upproaches [12. 8]. We also refer to, e.g. [7, 25, 22, 35, 19, 37, 10, 28), for replicating the "aha moment" of reasoning described in DeepSeek-R1. The core step of RI. for incentivizing reasoning ability is to use a rule-based reward, e.g., the correct answer of the math question, to provide training signal towards fitting the correctly generated answer by the model itself with reasoning traces. It has been observed that as the RL training progresses, the sequence length will be increased [8]. Apart from RL techniques, it has been shown that the simple SFT training pipeline with reasoning data distilled from strong reasoning models such as RI can also incentivize reasoning ability of LLMs; see, e.g.. [24, 33, 15, 30J. The length of SFT reasoning traces varies widely, ranging from 8k to 32k, and in some cases exceeding 100k. Training models on these long traces requires significant memory cost in storing activation values. The essential roles of RL and SFT for incentivizing reasoning ability are still under active discussion [36].
Memory-efficient training methods. Existing memory-efficient solution for training LLMs mainly focus on designing parameter-efficien or low-dimensional version of Adam [13, 20], significantly reducing the memory usage of storing the gradient and optimizer states. Typical algorithms include LoRA [9] and its variants, e.g., [5, 18]. GaLore [38], BAdam [23], etc. However, in the context of training reasoning models with long reasoning sequences, the dominant source of memory consump- tion arises from the BP process for storing intermediate activations. Gradient checkpointing technique [4] reduces the memory cost by recomputing the activation during the backward process. However, it still stores the full activation of the reforwarded layer and the full logits of the output layer; see Figure 1 for illustration. MsT [21] reduces the memory cost of SFT logits by iteratively processing mini-sequence. However, it still requires the full storage of layer activations during reforward and does not apply to RL objectives. BPT [17] reduces the activation memory in transformer layer by partitioning the forward process across the sequence dimension, yet it does not handle the logits memory cost in language modeling head layer. Sequence parallel approaches increase the maximum sequence length by partitioning sequence across GPUs [16, 11, 14], which requires access of multiple GPUs. MoNET [26] explores operator-level memory optimization, yet is specifically designed for vision models and cannot be applied to the training of LLM.
Main contributions. Motivated by the above observations, this work aims to address the memory issue occurred during the BP process. Our main contributions are summarized below.
(C.1) We propose memory-efficiett and exact backpropagation algorithm called StreamBP, which significantly reduces memory usage of storing the intermediate activation values when training LLMs ultra-long sequences, e.g., training reasoning models. StreamBP is based on a linear decomposition of the chain rule and involves nontrivial and intricate developments for the language modeling head and transformer layers. As result, StreamBP is compatible with common training objectives, including SFT, GRPO, and DPO. Moreover, by leveraging the causal structure of LLM, StreamBP is able to save computational FLOPs compared to the standard BP, achieving faster BP speed than the gradient checkpointing bascline. To support mulli-GGU training. we also develop distributed implementation of StreamBB with special attention to gradient and parameter communication, significantly improving training efficiency and broadening its applicability.
(C.2) Empirically, we measure the maximum sequence length and time cost of StreamBP under both single GPU and distribuetd training settings. Specifically, we measure the memory and time cost of BP in Section 4.1. StreamBP substantially scales up the maximum sequence length by 2.8 - 5.5x compared to the gradient checkpointing baseline across different model scales while achieving comparable or even less BP time. In Section 4.2, we verify the effectiveness of StreamBP under various training objectives, showing that it significantly increases maximum sequence length under the objective of SFT, GRPO, and DPO. In Section 4.3, we show that under Deepspeed distributed training scheme, StreamBP is able to achieve about - 5.6x larger sequence length than distributed gradient checkpointing.
2 Preliminary: Peak Memory Cost during Backpropagation
As a concrete example, we utilize PyTorch CUDA memory snapshot tool to record detailed GPU memory usage across 2 forward and backward processes of Qwen 3-4B model under the sequence length of 8192, where the gradient checkpointing technique is applied. The result is shown in Figure 1. We exclude the memory cost of optimizer states and focus solely on the BP process.
2

---

Layer
30G
FP32 logits grad
activations
Reforward
2nd backward
2nd forward
www.
200
backward
BF16 logits copy
153
1st forward
106
BF16 logits copy
Checkpointed layer activations
Model parameters
Figure 1: Memory profile during backpropagation of Qwen 3-4B with gradient checkpointing, visualized using PyToreh's memory profiler. Sequence length is set 8192. Batch size is Optimizer states are excluded as we focus on the BP process.
Peak memory cost. The peak memory cost occurs at the end of the 2nd forward pass. Apart from the 16GB allocated for BF16 parameters and gradients, approximately 14GB is used for storing intermediate computations. This includes BFI6 checkpointed layer inputs, BF16 logits, FP32 logits, and the gradient FP32 logits. Since the logits is dimension sequence length (T) x vocabulary size (C), their memory cost is independent of the model size and is subject to C' of the model class.
Second peak memory cost. The 2nd memory peak happens at the start of the 2nd backward, where the model reforward the checkpointed inputs of the last transformer layer to compute layer activations. The activation values will be temporarily stored for computing the gradient of the layer's parameter and inputs.
We remark that as the model size grows up, the 2nd peak memory will increase and become closer to the peak memory. We provide more detailed explanation of the memory profile in Appendix C.
3 Memory-Effiiint Exact Stream Backpropagation
In this section, introduce the design of the proposed algorithm, which computes exact gradient with reduced memory cost and floating point operations (FLOPs) during the BP process.
3.1
Main Idea
Consider a transformation happened during the model's forward pass fw Zout- where W is the weight associated with the transformation, fw (.) be any mapping inside the model, e.g., transformer layer or the language modeling head. By chain rule, the gradient of weight are given
by [vec(W] @vecl Zen) (ve )vec(Zau)
0L
aL
where L is the loss, vec(.) is the vectorized operator, and
@vec( Zour) (Ove)))ec(WW))) denotes the Jacobian matrix. During BP, when the gradient åL/dvec( Zoor) ready, gradient checkpointing method will reforward Zi through fw(-). and then calculate and store all the intermediate activations that are required for computing Ôvec( /Ove) and (rec(Zow)/Oveer Zm). where Òvec Lowi) /vvee will be used for calculating åL Ōvec Zom) ) for the preceding layer. As shown in Figure 1, the memory cost of storing these intermediate activation values can be huge.
To reduce the memory cost of these intermediate values, we introduce the stream backpropagation (S) Let vec(Zoat) = |vec 2())), vec v (ee((((D)) be any partition of vec(Zom). StreamBP is based on the following linear decomposition:
aL @vec(W))
/övec( Zour) T
aL
Ovec ((1) (vee(Z(()) @vec(W)
aL
Σ Em1
(1)
@vec(W))
Ovec Zout)
Övec z

---

By strategically partitioning vec(Zou., storing the intermediate activations required for comput- ing Övec ( (W)) can be significantly cheaper than storing those required for calculating Övec( Zom) / (W") and is often proportional to the partitioned chunk size. Motivated by this fact, StreamBP sequentially computes the decomposed components in (1) across all the partitions. and accumulates them in running sum, yielding the exact gradient. We refer to Appendix A. for quick grasp of StreamBP's idea when using linear transformation.
Next, we elaborate on applying StreamBP to concrete transformer LLM, which necessitates nontrivial and intricate developments.
3.2 StreamBP for Transformer LLMs
In this section, we apply StreamBP to significantly reduce the memory cost consumed by language modeling head and transformer layer during BP. Additionally, we will discuss that StreamBP requires even less computational FLOPSs compared to the standard BP with gradient checkpointing, enabling potential time acceleration over standard approaches.
3.2.1 StreamBP for Language Modeling Head: SFT, GRPO, and DPO
The language modeling head performs the following linear transformation:
H H im
head head
logits.
Here, Win head RdxC the weight of language modeling head, H RRxx is the hidden states output of the last transformer layer. The logits E IR' will be used to compute the objective function. C,d, d, T' are the vocabulary size, hidden dimension, and sequence length, respectively. As shown in Figure I, the logits and its gradient give rise to a huge memory consumption, due to the large vocabulary size and sequence length. the following, we analyze one-by-one how the memory of logits can be significantly reduced using StreamBP in the regime of SFT, GRPO, and DPO.
Supervised finetuning (SFT). The (un-normalized) "code-style" objective of SFT is given by
(logits Y) =
log softmax(logitss_)))
where E RT is the label vector. Importantly, each position's logits contribute to the objective independently. To perform StreamBP, the logits and label are evenly partitioned into D chunks across
the sequence dimension, i... { (logits y I...., D}
D} with logits(f) (f) (((TTI)/D)XC Then,
we sequentially accumulate the gradient across all partitions i = =. D as
Pin head (JSFr(logits: Y())
Wlssr(logits (9) (0 y())
9H
WWm_head
0H
where 9im_beae and 9и initialized from zero. The operator += denotes the in-place summation. logits()) and its gradient will be cleaned from memory once they have been used in (3). After the accumulation across all partitions, Ilm_head and 9и will be the exact gradient of the Wmm head and Ht, respectively. During this computation, StreamBP only stores logits (6) and its gradient sequentially for all i, which only costs 1/D memory compared to the original approach.
Group relative policy optimization (GRPO). The principle of StreamBP for GRPO aligns with SFT. For compact presentation, we display the detailed derivation in Appendix A.2.
Direct preference optimization (DPO). The objective DPO is given by
To ( 31 (y ( Bw.<e) log
Llogit) := -E log o (33) 1=1 log
Trer( 220, (a)
(Tee ( MI, (1,
(E(()
Here, logits (logits logits aed with logits, TTC a(.) is the sigmoid function. T is the sequence length of response y. To ease expression, we assume yoe and y have the same length without loss of generality. The policy's output determined by logits:
TV)(2, Mee) softmax(logits,.)))

---

Self Attention
MLP
[d/GjTT
M(O
repeac
K (cached;
Mgate Tx dgate
Partition loop
inn
Ql
Attention mask M: TxT
doo Hortzoo
Attention Operator (e.g. flash attention)
OUT
Hin Txd
Q: Txd (coched)
viil
( Hup
Hout Txd
repeai CX
VT: (dG))T
Hup dup
Figure 2: StreamBP for transformer layer (best view in color). The stored activations StreamBP is highlighted in orange. Unlike the gradient checkpointing approach that reforward H compute all the activatioins required for the backward of Hout, StreamBP computes the activations only for one partition of H time, which reduces the memory cost by large margin.
Unlike SFT and GRPO, DPO's objective cannot be divided into summation of losses on individual partitions, due to the existence of the non-linear log-sigmoid transformation. Fortunately, its gradient retains separable structure:
LLPO OW
(1-
al(t)
= -E Bw)
(5)
aW
StreamBP partitions logits across the sequence dimension flogits ()) = flogits logits (
Hi Rnd
1..... D} with logits E (TT/D)xC Based on the partition, performs the following accumula- tions for =1 ll..., D:
ae(I)
de(L)
e+= LET: E(t),
bead bead
E TET. OW ie
9H +- 3
(6)
LET. OH
where T, == {(i- - ((3-1)3 t i E Z}. After finishing the above accumulation, StreamBP performs the following in-place correction to compute the exact gradient:
Dem_head + (o(BE) 1m
9H - (a(3l) - 1)gH.
3.2.2 StreamBP for Transformer Layers: Attention and MLP
We now apply StreamBP to the transformer layer. To case presentation, we disregard components such normalization layer, multi-head mechanism, and residual connection without loss of generality.
transformer layer consists of consecutive transformations of attention and MLP:
Jam (Hm)): Hout = fMLP(O),
(8)
where Hi Rxx and Hom E R'xx the input and output of the transformer layer, respectively. The two transformations are given by
Attention:
Q H Wa, K HWn
V= HWWW
Q,K, V E Txxd
S = QK E RXT,
P= = S0) E RXT, -= PV e T'dd
MLP:
Hup = OW E IRTxs Hgate = =W gute E Txxx
Hout (Haae)) Hup) Woown E TRTxd
Computing oHoo HH requires storing Q, K, M,M,CR O, Hup. H Hgater Homt where M E RXT is the attention mask. which is required when using mini-batch training or techniques such as sliding window attention. The storage of S and P can be avoided using flash attention's 's tiling approach.
Define the partition of quantity E T''xd along the sequence dimension as { = I,... D}
with (10) E DT/ Let (+)) E (DTT/D)) be the concatenation of (H6))}_ along the sequence dimension. StreamBP for transformer layer is built upon the following observation:
5

---

Property 3.1 The computation of HH /oW only depends O6,, Q K(0, and V
To justify the above property, StreamBP sequentially performs the following partitioned attention and MLP for each chunk i:
Partitioned Attention:
Q Ho
k
H V H W. p()) softmax M()),, O - p(0
WR.
5(i) Q k(O)
Partitioned MLP:
H Wup
H
W
gate
H (o(H
H
W Wdown (TTD))xd
(10)
When the partitioned reforward finished, the activation for computing Övee (H /Vvec(W)) is stored memory. Then, we accumulate the gradient for = =I..... D, as
vec(gw) + (Hve) @v) Ovec(Hadl)
T
ØL
dvec (H
T
aL
vec(gn.) +
(11)
Övec Hm))
Ôvec (H
When the accumulation is finished, 9w and 9H, become the exact gradient of weight matrices W : (W,. Wa, We, Wap. Weate Wdown] and Hin respectively. Note that K and (+a) are needed for each partitioned forward. Hence, compute K and V at once and cache it during StreamBP of
the current layer. We remark that the partitioned attention is compatible with the flash attention. We illustrate the StreamBP for transformer layer in Figure 2.
Memory efficiency of BP. StreamBP only needs to store the following activation values: Q(1, K, V, MO, O6+), H H(() gale and H Note that when grouped query attention [2] is used with group size G. K and V only costs 1/G memory of Q. Consequently, StreamBP costs approximately 1/D memory for activation value compared to the standard BP.
Computational efficiency of StreamBP. For long sequence training, the most computational expen- sive operation involved in transformer layer is the calculation of the pre-attention score S. StreamBP reduces the FLOPs of the operation by approximately half. Specifically, the standard implementation
= QKT costs FLOPs of dl, while StreamBP only costs FLOPs of (1+D)T222
for computing
D
s()) = ((a) across partitions. The FLOPs reduction is since that StreamBP utilizes the causal structure of language models, which uses (*+))) rather than in the computation of S(9), For all the other operations, StreamBP across all partitions shares the same FLOPs the standard BP with checkpointing. Note that K and V are only computed once and get cached.
HBM overhead of StreamBP. Each partitioned gradient calculation in (1) requires loading model weight W from high bandwidth memory (HBM) register for computation, which induces additional overhead compared to the standard BP. Meanwhile, StreamBP reduces the HBM throughput of attention mask by approaximately half. The overhead directly depends on the number of partitions D. In Section 4.5, we empirically study how D) affects the BP time and memory cost.
3.3 Distributed StreamBP
Although StreamBP is naturally compatible with modern distributed training techniques such as distributed data parallel (DDP) and Deepspeed ZeRO, directly applying these techniques to StreamBP is inefficient, due redundant communications of gradients and parameters. We display the design of distributed StreamBP in Figure 6. The major designs contain gradient communication design and parameter communication design.
Gradient communication.
During the backward, when the local gradient buffer all processes are ready, a gradient averaging operation (i.e. all-reduce reduce-scatter) is performed across all the processes. As shown in (1), a parameter's gradient in StreamBP is ready until the accumulation oper- ation is finished. To avoid redundant gradient communication, the gradient averaging of distributed StreamBP is performed after the accumulation. this sense, distributed StreamBP will have the same gradient communication cost the standard BP.
Parameter communication. Parameter communication required when using model-parallel. For instance, ZeRO-3 partitions each parameter evenly across different GPUs, and gather the parameter

---

4.7x20l:
110k
1.7k
40k
440
5.7k
16k
GPU Memory (GB)
80
70
75
75
70
Buse wi cipt
50
65
70
Weights ARto (80688]
65-
30,5GB
55,0GB
-63.GBB
25k
75k
10k
20k
Sequence Length
Sequence Length
Sequence Length
(a) Qwen 3-8B
(b) Qwen 3-14B
(c) Qwen 3-32B (LoRA Grad.)
Figure 3: Peak BP memory cost measurement of Qwen 3-8B, 14B, and 32B models under different sequence lengths. Under the 80GB memory limit, StreamBP scales the maximum sequence by 2.8 5.5x larger compared to gradient checkpointing.
Sequence length
6k
9k
12k
15k
18k
21k
24k
27k
Baseline w/o ckpt
2.0
Baseline w/ ckpt
2.5
4.7
7.4
10.8
14.7
19.3
24.3
MsT
2.6
4.9
7.6
11.2
15.2
20.3
25.2
31.8
StreamBP
2.6
4.7
7.0
10.0
13.2
17.2
21.2
25.8
Acceleration over ckpt
4.4%
0.4%
6.0%
7.4%
10.5%
10.9%
12.9%
Table 1: BP Time cost (in seconds) under different sequence lengths. The result is based on Qwen 3-4B model and is averaged over 50 independent trials. The acceleration of StreamBP becomes more apparent as the sequence length grows up, corroborating our analysis in Section 3.2.2.
when the operators associated with is called. By (1), StreamBP requires the access of unpartitioned parameter for D times during the backward. Hence, naively implementing model-parallel for StreamBP costs (D 1)x all-gather operations. To reduce the communication cost, distributed StreamBP caches the weight locally until the accumulation of its gradient is finished, avoiding inducing additional parameter communication cost compared the standard BP. The design will introduce an additional memory cost of storing single transformer layer in each GPU.
4 Experiments
In this section, we evaluate StreamBP based on the Qwen model series. The result directly applies for any causal language models such as Llama, Mistral, and Gemma, since StreamBP is not restricted certain model class. Our evaluation mainly consists cost: 2) training cost; and 3) distributed training. All the experiments are conducted using 8800-80GB GPUs. The detailed setup is presented in Appendix D.
parts, including backpropogation
4.1 Backpropagation Cost Measure
Memory cost. We report the memory cost StreamBP and baseline (i.e., standard BP) with/without gradient checkpointing in Figure 3. Given the 80GB memory limit, StreamBP is able to increase the maximum sequence length by 2.8-5.5x and 23.4-36.3x compared to baseline with/without gradient checkpointing, respectively. Importantly. the memory cost of StreamBP is linearly related to the sequence length, as StreamBP does not store the full attention mask and is compatible with flash attention. Thus, StreamBP's sequence length scaling can be directly transferred to batch size scaling for accelerating the training. Note that the ratio differs across different model sizes, which attributes to distinct configurations of hidden dimensions.
Time cost. We compare the BP time cost of StreamBP with baseline approaches, and present result in Table 1. The result shows that StreamBP consistently exhibits faster BP time compared to the gradient checkpointing baseline across wide range of sequence lengths, and beats the oon--sequence training baseline MsT [21] more significantly. As the sequence length increases, the acceleration
7

---

4B
8B
14B
32B
Objective
Method
Full
LoRA
Full
LoRA
LoRA
LoRA
Baseline w/o ckpt
7.0
7.1
3.4
5.1
2.9
0.4
SFT
Baseline w/ ckpt
28.5
36.5
15.7
30.6
23.0
5.9
StreamBP
200.0
247.0
72.0
142.4
84.6
16.3
Baseline w/o ckpt
0.9
1.0
0.7
0.4
GRPO
Baseline w/ ckpt
8.5
12.1
9.5
6.9
StreamBP
18.2
27.5
16.5
10.1
Baseline w/o ckpt
3.1
3.5
2.5
1.4
0.2
DPO
Baseline w/ ckpt
18.7
32.2
25.5
18.5
4.4
StreamBP
57.2
100.2
60.9
36.9
7.8
Table 3: Maximum sequence length (in thousands) on a single A800-80GB GPU.
becomes more significant. The result aligns with our analysis in Section 3.2.2, where StreamBP reduces approximately half of computation in calculating attention scores, whose computational FLOPs is quadratically related to the sequence length. In Table 2, we show that StreamBP's memory- efficiency enables the usage of much larger batch size, which helps further accelerate the BP speed. The acceleration attributes to the reduced sample-wise HBM throughput of model loading and better leveraging of parallel processing. particular, StreamBP with batch size achieves even less per-sample BP time than the baseline without gradient checkpointing under batch size 1.
82%
Batch Size
1
4
16
GPU GPUMMe Memory (GB)
Baseline w/ ckpt
7.42
7.12
7.04
MsT
7.64
7.41
7.28
7.06
StreamBP 7.03
6.82
6.63
6.47
MrT
Table 2: Per-sample BP time cost of Qwen 3-4B model under different batch sizes. The sequence length is 9000. StreamBP achieves further accelera- tion by utilizing substantially larger batch size.
StramBP
A800 (59G8)
25k
100%
125k
Sequence Length
Figure 4: Memory cost comparison with MsT.
4.2 Training Cost Measure
Sequence length scaling. We present the maximum sequence length of StreamBP and baseline methods given single A800-80GB GPU in Table 3. For all the objectives, StreamBP significantly increases the maximum sequence length over baseline approaches, which justifies the efficiency of StreamBP in reducing the memory cost of transformer layer activations and logits. For SFT, even 32B model can achieve sequence length to 16k. Note that the sequence length scaling ratio can be equally transferred batch size scaling ratio for acceleration. For example, for the SFT of 8B model, StreamBP enables 72 R 4.5x larger batch size than the baseline with gradient checkpointing.
15.7
Comparison with long-sequence training baseline. We compare StreamBP with the long-sequence SET baseline MsT, and plot the peak memory training Qwen 3-8B LoRA model Figure 4. The result shows that StreamBP achieves approximately 1.7 larger sequence length than MsT. In terms of time-efficiency. StreamBP requires significantly less BP time compared to MsT, as shown in Table The acceleration becomes more significant the sequence length scales up.
4.3 Efficiency under Distributed Training
We measure the maximum sequence scaling and BP time under the Deepspeed ZeRO-2 SFT scheme, where the gradient and optimizer states are partitioned across different GPUs and each GPU processes its local data batch. The results are shown in Table and Table 5. Compared to baseline with gradient checkpointing, distributed StreamBP scales to a maximum sequence length approximately 5-5.6x larger and achieves noticeably faster BP speed.

---

# of GPUs
4
5
6
7
Baseline w/ checkpoint
18.6
20.5
22.3
23.1
23.5
23.7
Distributed StreamBP
92.4
112.0
121.5
128.7
131.4
131.8
Table 4: Maximum sequence length (in thousands) of Qwen 3-8B under ZeRO-2 training scheme.
Sequence length
6k
9k
12k
15k
18k
21k
Baseline checkpoint
8.4
9.6
12.2
18.5
22.9
26.3
Distributed StreamBP
6.4
8.3
10.8
14.8
17.2
20.7
Table 5: Per-sample BP time cost (in seconds) of Qwen 3-8B under ZeRO-2 training scheme.
4.4 Correctness Verification of StreamBP
We empirically measure the gradient difference of StreamBP and baseline methods to justify the cor- rectness of our implementation. Importantly, the float point operations has the following associativity issue due to numerical precision:
(a + b) Dc (b
(12)
where @ denotes floating-point addition. Therefore, is impossible for the gradient difference to be zero given the different computation order. To verify the correctness, we use the gradient computed by the pure FP32 BP of the standard BP with gradient checkpointing as ground truth, denoted as base32. Then, calculate its difference with the gradient 9; obtained by pure BF16 BP of StreamBP and baseline BP, respectively. Additionally, we include a StreamBP run under FP32 precision, denoted as stream32, to verify its numerical equivalence to base32. Specifically, detine the absolute error and relative error as
Erabs (g) Σ, base32 gil. (g) 1 Σ
baee32 gil
(13)
[g hane32 10 -10-101
n
Table reveals two key findings: 1) stream32's micro-scale deviations (Erats 10-9) validate mathematical equivalence to base32, and 2) stream16 1eam16 demonstrates <0.04% relative deviation from base16, indicating that StreamBP exhibits no precision loss compared to standard 1flot16 computation.
We further verify the correctness of StreamBP via model training and display the loss of SFT, DPO, and GRPO in Table 7. The loss of StreamBP closely align with the standard BP. The tiny difference due to different order of float operations discussed above.
4.5 Ablation Study and Additional Experiments
Effect of partition size. We fix the partition size of transformer layer to be 0.5k, 1k, 2k, and 5k, examining the memory and time cost across wide range of sequence lengths. The result is shown in Figure 5. Note that partition will not be performed if it is larger than the sequence length. When the sequence length relatively small, e.g., less than 10k, the BP times under different partition sizes are close. However, the sequence length scales up, using a partition size that is too small introduces substantial overhead. The overhead attributes to the additional HBM throughput on loading model weight from HBM to register for computation and repetitive kernel launches. Fortunately, as shown in Figure 5b, large partition size only introduces marginal additional memory cost. Hence, for long sequence training, one can
relatively large partition size maximize training efficiency.
Additional experiments. We provide more experiments in Appendix B and Appendix C.2. Here are the summarized results: 1) We present the sequence scaling using a single RTX3090-24GB GPU. The result shows that StreamBP is able to scale up the maximum sequence length to 15k, which is about 4.4x larger compared to gradient checkpointing. 2) We present the memory profile of StreamBP in Figure The profile demonstrates a substantial memory reduction in logits and activation values during layer reforwarding, which is consistent with our analysis in Section 3. 3)

---

Erabs
Enjll
Module
base16
stream16
stream32
base16
stream16
stream32
in_head
2.56e-7
2.64e-7
6.66e-9 1.75e-7
1.43%
1.47% 6.59%
0.04%
Transformer layers
2.46e-5
2.47e-5
6.63%
0.04%
Table Gradient precision analysis. StreamBP the numerically correct and has nearly identical gradient as the standard BP.
SFT (Capybara)
DPO (Ultrafeedback)
GRPO (TL;DR)
Step
Base
StreamBP
Diff.
Base
StreamBP
Diff.
Base
StreamBP
Diff.
5.782
5.782
+0.000
1.702
1.702
+0.000
0.013
0.013
+0.0000
100
4.924
4.926
+0.002
1.087
1.090
+0.002
0.011
0.011
-0.0002
200
4.313
4.309
-0.003
0.938
0.935
-0.003
0.005
0.004
-0.0004
300
3.990
3.993
+0.003
0.876
0.879
+0.003
0.006
0.006
+0.0000
400
3.791
3.795
+0.003
0.836
0.840
+0.004
0.005
0.005
+0.0001
500
3.667
3.663
-0.004
0.801
0.796
-0.004
0.004
0.004
-0.0001
Table 7: Evaluation loss of StreamBP and standard BP under different objectives.
We conduct experiments on Llama and Gemma architectures in Appendix B.2, demonstrating that StreamBP applicable to general transformer models.
50
0.5ix partition
partition
300
Zk partition sise
1k partition size
Memory (GB)
45
1k partiticn size
partiticn size
Time (s)
pariition size
5k partition size
200
40
100
35
0
10k Sequence Length
30k
40k
10k
20k
30k
40k
Sequence Length
(a) Time v.s. sequence length.
(b) Memory v.s. sequence length.
Figure 5: Time and memory costs of StreamBP on Qwen 3-8B with varying partition sizes.
5 Conclusion and Discussions on Limitations
In this work, we developed memory-efficient and exact BP method called StreamBP. Compared gradient checkpointing baseline, StreamBP requires significantly less memory cost and enjoys faster BP time by leveraging the causal structure L.MMs. StreamBP can be used for long sequence training of any transformer LLMs. which finds wide applications such training reasoning models. We also developed a communication--eficceen distributed StreamBP to support multi-GPU training.
Limitations. Currently, StreamBP does not support MoE or multimodal models. Nonetheless, these can be addressed with simple implementation extensions, as the underlying principle remains unchanged. Additionally, StreamBP's partition size have clear impact BP time. This overhead can mitigated by employing a fused backward operator, which significantly reduces the HBM throughput. We leave these directions for future work.
Broader impacts. provide memory-efficient and exact BP method long sequence training of LLMs. This is a technical contribution that does not yield explicit negative societal impacts.
10

---

Acknowledgments and Disclosure of Funding
We thank the Area Chair and the anonymous reviewers for their insightful comments and suggestions, which have significantly improved the quality of the manuscript.
Lei Zhao is supported in part by the Major Project of the National Natural Science Foundation of China (NSFC) under grant 72293582 and in part by the National Key R&D Program of China under grant 2023 122399 15202. Xiao Li is supported in part by the National Natural Science Foundation of China (NSFC) under grants 12201534, 12571330, and 12326608, in part by the 1+1+1 CUHK- (AHHKSZ)-GDSTT Joint Collaboration Fund under grant 2025A0505000044, and in part by the Shenzhen Science and Technology Program under grant RCY RCYX20221008090333000
References
[1] Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidi, Sam Altman, Shyamal Anadkat, et al. GPT-4 technical report. arXiv preprint arXiv:2303.00 087744000.77..... 2023.
[2] Joshua Ainslie, James Lee-Thhee Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebron, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head checkpoints. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pages 4895-4901, 2023.
[3] Sébastien Bubeck, Varun Chandrasekaran, Ronen Eldan, Johannes Gehrke, Eric Horvitz, Ece Kamar, Peter Lee, Yin Tat Lee, Yuanzhi Li, Scott Lundberg, et al. Sparks of artificial general intelligence: Early experiments with GPT-4. arXiv preprint arXiv:2303. /27/2, 2023.
[4] Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin, Training deep nets with sublinear memory cost. arXiv preprint arXiv:1644.0611 664406144. 2016.
[5] Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, and Luke Zettlemoyer. QLoRA: Efficient finetuning of quantized L.MSs. Advances Neural Information Processing Systems, 36, 2023.
[6] Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad AI-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The llama 3 herd of models. arXiv preprint arXiv:2407.2178. 2024.
[7] Hugging Face. Open R1: A fully open reproduction of deepseek-rl, January 2025.
[8] Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, et al. Deepseek-rl: Incentivizing reasoning capability in Ilms via reinforcement learning. arXiv preprint arXiv:2501.12945 2025.
[9] Edward Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang. Lu Wang. Weizhu Chen, al. LoRA: Low-rank adaptation of large language models. ICLR, 1(2):3, 2022.
[10] Jingcheng Hu, Yinmin Zhang, Qi Han, Daxin Jiang, Xiangyu Zhang, and Heung- Yeung Shum. Open-reasoner-zero:: An open source approach to scaling up reinforcement learning the base model. arXiv preprint arXiv:2503.242900 2025.
[11] Sam Ade Jacobs, Masahiro Tanaka, Chengming Zhang, Minjia Zhang, Shuaiwen Leon Song, Samyam Rajbhandari, and Yuxiong He. Deepspeed ulysses: System optimizations for enabling training of extreme long sequence transformer models. arXiv preprint arXiv:2309. 2023.
[12] Aaron Jaech, Adam Kalai, Adam Lerer, Adam Richardson, Ahmed El- El-Kishky, Aiden Low, Alec Helyar, Aleksander Madry, Alex Beutel, Alex Cameey, al. Openai ol system card. arXiv preprint arXiv:2412.1672D 2024.
[13] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.

---

[14] Vijay Anand Korthikanti, Jared Casper, Sangkug Lym, Lawrence McAfee, Michael Ander- sch, Mohammad Shoeybi, and Bryan Catanzaro. Reducing activation recomputation in large transformer models. Proceedings of Machine Learning and Systems, 5:341-355 2023.
[15] Dacheng Li, Shiyi Cao, Tyler Griggs, Shu Liu, Xiangxi Mo, Eric Tang, Sumanth Hegde, Kourosh Hakhamaneshi, Shishir Patil, Matei Zaharia, et al. Llms can easily learn to reason from demonstrations structure, not content, is what matters! arXiv preprint arXiv:2502.07374 2025.
[16] Shenggui Li, Fuzhao Xue, Chaitanya Baranwal, Yongbin Li, and Yang You. Sequence paral- lelism: Long sequence training from system perspective. arXiv preprint 1rXiv:2105.13120. 2021.
[17] Hao Liu and Pieter Abbeel. Blockwise parallel transformers for large context models. Advances in neural information processing systems, 36:8828-8844 2023.
[18] Shih- Shih-Yang Liu, Chien-Yi Wang. Hongxu Yin, Pavlo Molchanov. Yu-Chiang Frank Wang, Kwang-Ting Ting Cheng, and Min-Hung Chen. Dora: Weight-decomposed low-rank adaptation. In Forty-first International Conference on Machine Learning, 2024.
[19] Zichen Liu, Changyu Chen, Wenjun Li, Penghui Qi, Tianyu Pang, Chao Du, Wee Sun Lee, and Min Lin. Understanding r1 -1-zero-like training: A critical perspective. arXiv preprint arXiv:2503.2078. 2025.
[20] Nya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference Learning Representations.
[21] Cheng Luo, Jiawei Zhao, Zhuoming Chen, Beidi Chen, and Anima Anandkumar. Mini-sequence transformer: Optimizing intermediate memory for long sequences training. arXiv preprint arXiv:2407.15892, 2024.
[22] Michael Luo, Sijun Tan, Justin Wong, Xiaoxiang Shi, William Y. Tang, Manan Roongta, Colin Cai, Jeffrey Luo, Li Erran Li, Raluca Ada Popa, and lon Stoica. DeepScaleR: Surpassing ol-preview with a 1.5b model by scaling rl, 2025. Notion Blog.
[23] Qijun Luo, Hengxu Yu, and Xiao Li. BAdam: A memory efficient full parameter optimization method for large language models. Advances in Neural Information Processing Systems, 37:24926-24958, 2024.
[24] Niklas Muennighoff, Zitong Yang, Weijia Shi, Xiang Lisa Li, Fei-Fei, Hannaneh Hajishirzi, Luke Zettlemoyer, Percy Liang, Emmanuel Candès, and Tatsunori Hashimoto. sl: Simple test-time scaling. arXiv preprint arXiv:2501.19399, 2025.
[25] Jayy Pan, Junjie Zhang, Xingyao Wang, Lifan Yuan, Hao Peng, and Alane Suhr. Tinyzero. 2025. Accessed: 2025-01-24.
[26] Aashaka Shah, Chao- Yuao Wu, Jayashree Mohan, Vijay Chidambaram, and Philipp Krähenbühl. Memory optimization for deep networks. arXiv preprint arXiv:2010. 14501, 2020.
[27] Leandro von Werra, Younes Belkada, Lewis Tunstall, Edward Beeching, Tristan Thrush, Nathan Lambert, Shengyi Huang, Kashif Rasul, and Quentin Gallouédec. Trl: Transformer reinforce- ment learning. https: / / thub. hwtm:/uuginggaace..c.c
2020.
[28] Yiping Wang, Qing Yang. Zhiyuan Zeng, Liliang Ren, Lucas Liu, Baolin Peng, Hao Cheng, Xuchai He, Kuan Wang, Jianfeng Gao, et al. Reinforcement learning for reasoning in large language models with one training example. arXiv preprint arXiv:2504.20577 2025.
[29] Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia. Ed Chi, Quoc V Le, Denny Zhou, et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in neural information processing systems, 35:24824-24883 2022.
[30] Liang Wen, Yunke Cai, Fenrui Xiao, Xin He, Qi An, Zhenyu Duan, Yimin Du, Junchen Liu, Lifu Tang, Xiaowei Lx, et al. Light-RI: Curriculum sft, dpo and rl for long cot from scratch and beyond. arXiv preprint 20X1v22503.10460. 2025.
12

---

[31] Thomas Wolf. Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer. Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Transformers: State-of-the- art natural language processing. In Proceedings of the 2020 Conference Empirical Methods in Natural Language Processing: System Demonstrations, pages 38-45, Online, October 2020.
Association for Computational Linguistics.
[32] An Yang. Baosong Yang. Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chengyuan Li, Dayiheng Liu, Fei Huang, Haoran Wei, et al. Qwen2. 5 technical report. arXiv preprint WrX1::442.151155 2024.
[33] Yixin Ye, Zhen Huang, Yang Xiao, Ethan Chern, Shijie Xia, and Pengfei Liu, Limo: Less is more for reasoning. arXiv preprint UrXiv:2502.03387 2025.
[34] Edward Yeo, Yuxuan Tong, Morry Niu, Graham Neubig, and Xiang Yue. Demystifying long chain-of-thought reasoning llms. arXiv preprint arXiv:2502.03373, 2025.
[35] Qiying Yu, Zheng Zhang, Ruofei Zhu, Yufeng Yuan, Xiaochen Zuo, Yue, Tiantian Fan, Gaohong Liu, Lingjun Liu, Xin Liu, et al. Dapo: An open-source Ilm reinforcement learning system at scale. arXiv preprint ar 2rXiv 2025.
[36] Yang Yue, Zhiqi Chen, Rui Lu, Andrew Zhao, Zhaokai Wang. Shiji Song, and Gao Huang. Does reinforcement learning really incentivize reasoning capacity in ilms beyond the base model? arXiv preprini arXiv:2504. 13837, 2025.
[37] Weihao Zeng, Yuzhen Huang, Qian Liu, Wei Liu, Keqing He, Zejun Ma, and Junxian He. Simplerl-zoo: Investigating and taming zero reinforcement learning for open base models in the wild. arXiv preprint arXiv:2503. 18892, 2025.
[38] Jiawei Zhao, Zhenyu Zhang. Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. Galore: Memory-effictte Ilm training by gradient low-rank projection, arXiv preprint a0Xiv:2403.03357,, 2024.
13

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The claims in abstract and introduction supported by the memory and FLOPs analysis in Section and numerical results in Section
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We include the discussion in Section
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
14

---

Justification:
This paper does not prove any theoretical result.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We provide the full code and environment setup for reproducing the results of the paper.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
15

---

Answer: [Yes]
Justification:
We provide the full code and data simulation script in the supplemental
material.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips. ec/
oolloooioo icyy for more details.
While encourage the release of code and data, understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: / //nips. cc/
Seriiiees/ooSSubiiii swisssolicccyy for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We provide the detailed experiment setting in Appendix D.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer:
[Yes]
Justification:
The reported time cost is based on average 50 independent measurements.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
16

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: We report the computational resources used for all the experiments Section and Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (c.g., preliminary or failed experiments that didn't 't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https : http///eurip...eee wwwpuurip.cccppupaeeeeee httpppulit/tiiituiiiiiiiiiooeeee
Answer: [Yes]
Justification:
We have read the code of ethics and make sure that this research aligns with
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer:
[Yes]
Justification: include the discussion of broader impacts in Section
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
17

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The paper poses no risks mentioned in the question.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We have cited the assets used in this paper, e.g., code library and models.
Guidelines:
The answer NA means that the paper does
existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
18

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Question:
Answer: [Yes]
Justification:
We provide detailed guidelines on the usage of our code base.
Guidelines:
The answer NA means that the paper does release new assets.
Researchers should communicate the details of the datasettcooeemmodll as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: We do not conduct crowdsourcing experiments and research with human subjects in this project.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing experiments nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such the institution conducting the review.
19

---

16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
Justification:
The paper does not use LLMs for important, original, or non- non-standard compo- nents of the core method.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy (https: //tter/iss...///eeee weeeurips.cc/ooofeeeeeeeeeeeeeee /t////ffeeeneee2222222 20LL/LLM for what should or should be described.
20

---

Contents
Introduction
2
Preliminary: Peak Memory Cost during Backpropagation
3
Memory-Ffficient Exact Stream Backpropagation
3.1
Main Idea
3.2 StreamBP for Transformer LLMs
3.2.1
StreamBP for Language Modeling Head: SFT, GRPO, and DPO
3.2.2
StreamBP for Transformer Layers: Attention and MLP
3.3
Distributed StreamBP
Experiments
4.1
Backpropagation Cost Measure
4.2
Training Cost Measure
8
4.3
Efficiency under Distributed Training
4.4
Correctness Verification of StreamBP
4.5
Ablation Study and Additional Experiments
Conclusion and Discussions on Limitations
10
A
Additional Details of Stream Backpropagation
22
A.1 StreamBP for Linear Transformation
22
A.2 StreamBP for GRPO
22
A.3 Distributed StreamBP
23
Additional Experiments
25
B.1 Sequence Scaling on a Single RTX3090-24GB
25
B.2 StreamBP for Broader Architectures
25
Analysis of Memory Profile
25
C.1 Memory Profile Explanation of Gradient Checkpointing.
25
C.2 Memory Profile of StreamBP
26
Experiment Setup
26
21

---

A Additional Details of Stream Backpropagation
A.1 StreamBP for Linear Transformation
Given the input X = X2.... x E R Px and weight matrices W1 E Rmxx W2 E Rnxn, Consider the following two linear transformations:
Y Y1, ......... Ybl = XW1 E ER
Z Zi.... Zpl == YW2 E Px
For the ease of expression, define dM = aL the gradient of quantity M with respect to the objective L, which can be arbitrary objective. T gradient Y and weight matrices are given by
dY (d dW, x dW₂ = YT2=Y((D))) (dZ).
The standard approach of calculating dW1 and dW2 requires storing the intermediate value Y and dY. Based on (1), the above expression can be written as
D dY; = ( dW1 Σ (dr)), dW2 = X (dZ.). 1=1
D
Hence, one can sequentially compute Y; and dY, accumulate X, dY; and Y, dZ,. then drop the Y; and d',, Compared to the standard gradient computation, this approach effectively reduces the memory cost of intermediate values to 1/D without introducing additional computational cost. In practice, to better utilize the parallel computation and reduce the HBM load, one can process X in chunk-wise rather than sample-wise manner. Below, we demonstrate the memory efficiency of StreamBP with numerical experiment on the linear example.
Experiment on linear example. We empirically measure the memory cost of StreamBP and compare it with standard BP, based on concrete linear example. Specifically, let X E RR10 16384 W1. W2 E g16984x16384 L = Σj j.A Zj StreamBP partitions X as (X (2 = 1,... D} with
( E /D10 We present the memory cost and time cost of standard BP and StreamBP in Table 8.
StreamBP
Standard BP
20
D 100
D 200
D 500
Peak memory (GB) Intermediate memory (GB)
36.01 25.15
13.25
12.47
12.20
12.07
11.99
2.39
1.61
1.34
1.21
1.13
Time cost (seconds)
14.48
14.50
14.70
14.79
15.45
18.42
Table 8: Memory and time cost of standard BP and StreamBP under different partition number.
Compared to standard BP, StreamBP with D 20 reduces memory cost by 63.2% with almost time overhead. As D increases, the intermediate memory cost is further reduced.
A.2 StreamBP for GRPO
The GRPO's objective is given by:
CaRrollogito := (o
Tc
min
To () (j
A clip
To(J,() Tola(it)
- 6, 6,1+6) Asi - log
To 1) Teer Ther(j.
G
T
1=1
(14)
To ease presentation, the above equation we omit the compensation term GRPO's KI. divergence term without loss of generality, as StreamBP directly applies to the KL divergence term. Here, is
22

---

the prompt sequence, Dg is the dataset of prompt, and is j-th response sequence with respect to 9. G and To are the group size and the length of response. For the ease of expression, we assume all the responses are of the same length without loss of generality. A,, is the estimated advantage. ng. Vold- and Tee are the target policy, old policy and reference policy, respectively.
Here, logits - {logits,,, n,> logits nau? and reference policy, respectively, with logits, E I The policy's output is determined by logits, i.c.,
Mau? logits,.,,) contains logits generated by target policy, old policy,
T(j, () TO0 0j,cx) voftmax(logiiss,,...))))
Note that each f(jtt (14) contributes to the objective independently and only depends on logits,_j,t,>> which enables us to perform StreamBP along the sequence dimension. Specifically, let us
T. logits 1. 1...., D} with logits(" (t) E RGx(TT/D)xC, Define the objective of the sequence partition as
partition the logits along the sequence dimension as {logits()) := (logits
Fou" logits (7)
Ric Hi
GRPO (logits (1) ))= = 1B (g D.. -
(6)
[d
T
1 G DIET ((.tt)
(15)
where T: == 11 <t Io t 2} denotes the sequence range of the partition. We have LGRPO (logits) = ΣP' LORP (logits ). Then, similar to SFT, StreamBP sequentially performs the following accumulation for i -I....," D:
aL
(1(() o (l ())
9in_beat
9H +=
(16)
OW hmadddd
0H
A.3 Distributed StreamBP
StreamBP is naturally compatible with distributed training techniques such Deepspeed ZeRO. However, to ensure the efficiency, the gradient communication and parameter communication needs be customized. The design of distributed StreamBP is shown in Figure 6.
Gradient communication design. Gradient communication happens in almost all distributed training. During the backward, when the local gradient buffer of all processes are ready, a gradient averaging operation (i.c. all-reduce or reduce-scatter) is performed across all the processes. As shown in (1), parameter's 's gradient in StreamBP is ready until the accumulation operation is finished. To avoid redundant gradient communication, the gradient averaging of distributed StreamBP is performed after the accumulation. In this sense, distributed StreamBP will have the same gradient communication cost as the standard BP.
Parameter communication design. Parameter communication is required when using model- parallel. For instance, ZeRO-3 partitions each parameter evenly across different GPUs, and gather the parameter when the operators associated with it is called. By (1), StreamBP requires the access of unpartitioned parameter for D times during the backward. Hence, naively implementing model- parallel for StreamBP costs (D 1)x all-gather operations. To reduce the communication cost, distributed StreamBP caches the weight locally until the accumulation of its gradient is finished, avoiding inducing additional parameter communication cost compared to the standard BP. This design will introduce an additional memory cost of storing a single transformer layer in each GPU.
Efficiency of distributed SPreamBP. The choice of batch size significantly affects the communication cost. For example, when using Deepspeed ZeRO-2, using batch size 1 with gradient accumulation step K leads to approximately K times heavier backward communication cost than using batch size K with gradient accumulation step 1. StreamBP enables the usage of much larger batch size compared to the standard BP, which significantly reduces the communication cost. We remark ZcRO-2 reduces the above overhead by overlapping the communication and computation.
23

---

Layer
GPU 1
Partition losp
Transformer layer
Hin (lonal batch)
Hout (local batch)
Acsumulate dW and difn acroes parti
Reduoe-scatter
Parameter gathering (for model-paralee)
Parameter partition modal-parallel)
No gradient and parameter communication during SPreamBP of the layer.
GPU 2
Portition loop
doo) uazann
Transformer layer
Hin (local batch)
Mout (local batch)
Accumullte dW and dHin acroas
Figure 6: Design of distributed StreamBP. When backward through a transformer layer, its parameter gathered beforehand. During the StreamBP of the layer, no gradient or parameter communication is fired. When BP of the layer is finished, the gradient will be reduced across layers and the paramater will be sharded across GPUs.
24

---

B Additional Experiments
B.1 Sequence Scaling on Single RTX3090-24GB
We present the memory cost of training Qwen 3-8B LoRA model in Figure 7. Under the 24GB memory budget, StreamBP allows sequence scaling up to 15k, which 4.4x larger than the gradient checkpointing baseline.
GPU Memory (GB)
24
22
20
Rase w/ dept
dapt
18-
Stream HP
Whights Grads
3050 (24GB)
16
15.9GB
0
2k
7k
10k
12k
15k
Sequence Length
Figure 7: Peak BP memory cost measurement of Qwen 3-8B with LoRA (rank=32) single RTX3090-24GB GPU
B.2
StreamBP for Broader Architectures
In this section, we display the BP memory cost of Llama 3.1-8B and Gemma 3-1B. The results shown in Table 9.
Model
Method
3k
6k
9k
18k
Base w/o ckpt
54.2
Llama 3.1-8B
Base w/ ckpt
35.8
41.7
47.6
65.6
StreamBP
31.8
32.7
33.6
37.2
Base w/o ckpt
23.1
42.8
63.4
Gemma 3-1B
Base w/ ckpt
14.3
24.8
35.3
67.2
StreamBP
4.7
5.0
5,4
7.9
Table 9; BP Memory cost (GB) of Llama 3.1-8B and Gemma 3-1B under different sequence lengths.
Clearly, the memory cost of StreamBB is significantly less than standard BP; the memory gap increases the sequence length scales up, which verifies the efficiency of StreamBP different model architectures. Notably, Gemma 3-1B costs even more memory than Llama 3.1-8B for standard BP under sequence length 18k. This is because Gemma model has 2x larger vocabulary size than Llama 3.1, and hence doubles the memory cost of logits. In comparison, StreamBP only stores partitioned logits and its memory cost dominated by the model weight and model gradient rather than activation values.
C Analysis of Memory Profile
C.1 Memory Profile Explanation of Gradient Checkpointing
We provide detailed explanation of Figure 1. First, the "Model parameters" part stores the BF16 model parameters and constantly occupy nearly 8GB memory. Then, during the 1st forward process, gradient checkpointing gradually stores the checkpointed layer inputs, and hence it gradually consumes more memory. the end of the 1st forward and the beginning of the 1st backward, the FP32 logits and its gradient are computed, which suddenly consumes huge amount of memory due to the very large vocabulary size of Qwen 3. This huge memory occupation released after calculating the gradient of Im head (tied with embedding) as shown the brown rectangle above the "BF16 logits
25

---

copy". During the 1st backward, BP calculates and stores the gradients of all the parameters and the memory consumption continues to increase. One interesting phenomenon is that the memory usage of the checkpointed layer activations enclosed in the yellow triangle goes to decrease during the 1st backward process. This because that the stored corresponding activations will be deleted once the current weights' gradients are computed, as they are no longer needed. Another interesting observation is that there are some trianglc-shaped bumps during the 1st backward, for which we draw subfigure to interpret during the 2nd backward process. We explicitly show the reforward process of gradient checkpointing. and show the reforwarded layer activations that will be deleted once the corresponding gradients computed, yielding Htangle-shapee bump.
An additional remark on Figure is in order. As shown in Figure 1, the current implementation of the "Transformers" package stores an additional "BF16 logits copy" and an additional tied embedding and "Im head" gradient in the 2nd backward process, which may optimized as they are not used.
C.2 Memory Profile of StreamBP
The memory profile of StreamBP is shown in Figure 8. Compared to gradient checkpointing's profile in Figure 1, StreamBP reduces the peak memory greatly by partitioning the logits across the sequence dimension. The second peak memory in reforwarded activation values is also greatly reduced, which now only marginally higher than the storage of the key and value states.
Partitioned activation value
Partitioned logits & its grad
Cached K, V
18G
16G
14G
12G
10G
8G
6G
4G
2G
0
Figure 8: Memory profile with StreamBP under the same settings as Figure 1. The partition size of logits and transformer layer are set to 100 and 500, respectively.
D Experiment Setup
We implement our algorithm using Huggingface transformers library [31]. The detailed experiment setup is presented below.
Backpropogation cost. We decouple the optimization, focusing purely on the time and memory cost during the BP. This result serves as the minimum requirement on training a language model under a given sequence length. In particular, under batch size I, we measure the BP memory cost of Qwen 3-8B, 14B and 32B models using a single A800-80GB GPU. For 32B model, we inject rank-32 LoRA adapters and only calculate the gradient of the adapters, as a single A800 GPU cannot store the model and the full gradient simultaneously. We adopt the BF16 data type for storing model weight and gradient. The partition size of language modeling head is set to 100 for all the experiments. The partition size for transformer layer is set to 500 for maximum sequence length measurement, and is set to T/3 for time measurement.
Training cost. We measure the maximum sequence length of training 4B, 8B, 14B, and 32B models using the objective of SFT, GRPO, and DPO, respectively. Our implementation built upon HuggingFace TRL library [27]. We adopt pure BF16 Adam optimizer in full training and
26

---

mnixed-precisioe training in rank-32 LoRA model training. When using LoRA mode, there is no need store the reference model DPO and GRPO, as one can disable the adapter in training model to recover the reference model. The hatch size is set to I except for GRPO, where the group size is set 8. We also compare the memory cost of StreamBP with long-sequence training baseline method MsT to demonstrate the effectiveness of StreamBP. The partition size of language modeling head is set to 100 for all the experiments. The partition size for transformer layer is set to 500 for maximum sequence length measurement, and is set to T/3 for time measurement. All the experiments are conducted single A800-80GB GPU.
Distributed training. Under the communication-efien efficient design in Appendix A.3, we measure the maximum sequence length and time cost of distributed StreamBP under Deepspeed ZeRO-2 training paradigm, where the gradient and optimizer states are partitioned across GPUs. Our evaluation is conducted using a single-node server with A800-80GB GPUs connected by NvL.ink.
27


Reviews:
[{"review_id": "I4a4MD24I4", "reviewer": "Reviewer_ucaR", "summary": "This paper introduces StreamBP, a memory-efficient and exact back-propagation method for long-sequence training of LLMs. StreamBP uses layer-wise linear decomposition of the chain rule along the sequence, resulting in reduced memory cost of activation and logits. Compared to gradient checkpointing, StreamBP is able to support 2.8$\\times$ longer maximum sequence length.", "questions": "As listed in Weaknesses above. It would be great if the authors could elaborate on the above 3 points.", "limitations": "The current evaluation results are presented without any variance metrics (e.g. across seeds or runs), which weakens the confidence in the claim that StreamBP achieves.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\n- The paper intends to tackle a timely and important topic of reducing memory usage for long-sequence training of LLM;\n- The proposed method of streaming partial gradients to avoid storing full activations while keeping gradients exact (instead of approximating) seems novel and effective;\n- The empirical results demonstrate reduced activation memory usage on standard language modelling objectives e.g. SFT, DPO, and GDPO.\n\nWeaknesses:\n\n- Optimizing memory consumption for activation and logits is important, and the paper show effectiveness of streamBP towards this objective. However, in practice, it is not sufficient to consider only activation and logits for scaling large models and long sequences. As shown Fig. 3, model weights takes significant memory especially when model grows bigger. This should be explicitly discussed. One suggestion would be to discuss how StreamBP can complement weight, gradient, optimiser memory reduction methods such as GaLore and BAdam, to address the full memory bottleneck in large-model, long-context training.\n- The authors mentioned sequence parallel approaches increase the maximum sequence length, but there lacks further discussion or comparison with such methods. A short discussion of trade-offs and complementarity would strengthen the paper and help practitioners understand when StreamBP is preferable or combinable with existing parallelism strategies.\n- Current evaluation is on Qwen only, it would be good to see results with other SOTA LLMs to confirm StreamBP generalise well with different architectures.", "was_revised": true, "final_justification": "My concerns have been sufficiently addressed in authors' rebuttal.", "rebuttal": "We thank the reviewer for the valuable comments. We now address your concerns below.\n\n**How StreamBP complements existing memory-efficient methods for training LLM.** Thank you for the suggestion. The memory cost of LLM training mainly consists of storing activation values, model weight, and optimizer states (including model gradient). StreamBP reduces the memory cost of activation values during BP, and can be directly combined with existing memory-efficient training methods that reduce model memory (e.g., quantization methods, QLoRA) and optimizer states memory (e.g., LoRA, Galore, BAdam) for further reducing the memory bottleneck caused by model weight and optimizer states, respectively. Since StreamBP is an exact BP approach, combining it with other memory-efficient methods will not result in any performance degradation of the method. In our manuscript, we have provided training results for combining StreamBP and LoRA in Table 3. We will explicitly state that StreamBP can be directly applied to existing memory-efficient methods such as QLoRA, LoRA, Galore, BAdam in the next version of our manuscript.  \n\n**Comparison between StreamBP and sequence parallel.** Sequence parallel (SP) is specifically designed for distributed training with multiple available GPUs, while StreamBP is applicable for both single GPU and multi-GPUs settings. Specifically, SP partitions the activation across different GPUs, where each GPU receives one activation partition evaluated on the same input data. Consequently, SP requires a sufficient number of GPUs to scale to long sequences. Additionally, as different segments of the activation are stored in different GPUs, SP requires additional communication overheads between multiple GPUs. By contrast, StreamBP sequentially processes all partitions on the same GPU, storing only one partition at a time. This design allows StreamBP to facilitate long sequence training even in *a single GPU* setting, while also allowing further scaling of sequence length when multiple GPUs are available. We believe that this clearly clarifies the fundamental difference between StreamBP and SP. In the next version, we will add this discussion about these two approaches. \n\n**Evaluation on different model architectures.** We thank the reviewer for the suggestion. As StreamBP is an algorithm that linearly decomposes the chain rule during gradient computation in a transformer architecture, it is not tailored to any specific model architectures. Hence, it can be applied to any transformer LLMs. In the following tables, we present the BP memory cost of Llama and Gemma models using StreamBP and standard BP.\n\n| Sequence Length     |  3k    |   6k    |   9k    |  18k    |\n|---------------------|--------|---------|---------|---------|\n| StreamBP            | 31.81  | 32.66   | 33.62   | 37.20   |\n| Base w/ ckpt        | 35.79  | 41.69   | 47.61   | 65.59   |\n| Base w/o ckpt       | 54.20  | -       | -       | -       |\n\n**Table:** BP memory cost (GB) of Llama 3.1-8B.\n\n| Sequence Length     |  3k    |   6k    |   9k    |  18k    |\n|---------------------|--------|---------|---------|---------|\n| StreamBP            | 4.69   | 4.95    | 5.37    | 7.85    |\n| Base w/ ckpt        | 14.28  | 24.79   | 35.33   | 67.16   |\n| Base w/o ckpt       | 23.08  | 42.80   | 63.37   | -       |\n\n**Table:** BP memory cost (GB) of Gemma 3-1B.\n\nClearly, the memory cost of StreamBP is significantly less than standard BP; the memory gap increases as the sequence length scales up, which verifies the efficiency of StreamBP on different model architectures. Notably, Gemma 3-1B costs even more memory than Llama 3.1-8B for standard BP under sequence length 18k. This is because Gemma 3 model has 2$\\times$ larger vocabulary size than Llama 3.1, and hence doubles the memory cost of logits. In comparison, StreamBP only stores partitioned logits and its memory cost is dominated by the model weight and model gradient rather than activation values.\n\n**The evaluation results are presented without variance metrics.** Given the fixed model, sequence length, and partition size, the memory cost of StreamBP will be a constant and does not vary across different seeds and runs.\n\nWe hope that our response is satisfactory and addresses all your concerns. If there are any additional questions and concerns, please feel free to let us know during the author-reviewer discussion period. We are more than happy to provide further clarifications."}, {"review_id": "ZmxuNZvBsV", "reviewer": "Reviewer_ub3z", "summary": "This paper is about a technique to reduce memory footprint when gradient checkpointing is used.\nThe main idea is to split values in each feedforward layer into chunks, so that the entire values do not need to be recomputed in the re-forward step of gradient checkpointing.\nThe main target for the reduction technique suggested is the Transformer architecture, and various splitting strategies are suggested for each component of the architecture, including objective functions from SFT, GRPO, DPO and Transformer componenets like attention, MLP.\nDue to the splitting strategies adapted to each component, it can benefit in computation efficiency due to reduced sizes of matrices used in computing causal operators like attention.\nFrom evaluations, it is empirically proven that in terms of memory cost the StreamBP method outperforms baselines without checkpointing by 23-36x and those with checkpointing by 2.8-5.5x.\nAlso in terms of time cost, the proposed method generally outperformed the plain checkpointing technique, especially when the length of context gets longer (likely due to reduced computation in attention).", "questions": "* To my assumption in exchange of reduced memory cost the StreamBP method needs to do computation multiple times sequentially, i.e. for the number of chunks/blocks. However actually it resulted in less time required for training, is it solely due to partitioning of attention layers while considering causal relationship? If so, I want to see the speed degradation in other layers vs. speed improvement in attention layers.", "limitations": "It would be more nice to mention if there are cases that this technique cannot be used.", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 4, "originality": 3, "strengths_and_weaknesses": "Strengths\n* Good overall performance improvement\n* Even with more fine-grained checkpointing, it is actually performant in terms of time cost.\n\nWeaknesses\n* Comparison with possibly similar prior work like MoNET, Block Parallel Transformer would make the statement stronger.", "was_revised": true, "final_justification": "My concerns are resolved in the author response and I changed the score. Thank authors for the clarification.", "rebuttal": "We thank the reviewer for the valuable comments. We now address your concerns and questions below.\n\n**Comparison with MoNET [1]**\n* StreamBP and MoNET are designed for different tasks. In particular, MoNET is primarily used for vision tasks, and is specifically optimized for convolution-heavy operators, e.g., batch normalization and max-pooling. In comparison, StreamBP focuses on reducing the activation memory cost of the training of transformer LLMs, where the activation memory can usually be much higher. \n\n* The core methodology between StreamBP and MoNET are essentially different. MoNET dynamically determines which quantities to checkpoint for minimizing the memory consumption. By contrast, StreamBP adopts the layer-wise checkpointing strategy and uses partitioned forward and backward of a single layer for reducing the memory cost.\n\n**Comparison with Block Parallel Transformer (BPT) [2]** \n\n* BPT only reduces the activation memory in transformer layer, while StreamBP is applicable to general mapping, e.g., language modeling head, which further reduces the memory consumption by a large margin. By applying StreamBP on the language modeling head, the logits memory is significantly reduced by more than 90%; see Figure 8. In certain setting (such as Figure 1), the memory consumption of the language modeling head can be huge, and hence reducing the memory cost of this part could be important. \n* StreamBP reduces computational FLOPs by leveraging the causal structure of LLM. In particular, it saves approximately half computation of the pre-attention score; see lines 188-195. In comparison, BPT still computes the full pre-attention score. \n* StreamBP is compatible with most of the fused attention kernels, e.g., flash attention, and thereby has much less HBM throughput than BPT.\n\nWe note that BPT is implemented using JAX. It is highly nontrivial to reproduce the exact BPT implementation in PyTorch for an experiment comparison. We will provide further discussions on MoNET and BPT in the next version of our manuscript.\n\n**Speed degradation of applying StreamBP in layers other than attention layer.** The BP time acceleration of StreamBP comes from the reduced computational FLOPs in attention score calculation and the reduced HBM throughput of attention mask. To see StreamBP's effect on BP speed degradation, we apply StreamBP on the language modeling head only, and compare it to the gradient checkpointing baseline. We perform 2 rounds of BP on a single A800 GPU and summarize the results in the following table: \n\n| Partition   Size   | 100   | 200   | 500   | 1000  | 3000  | \\| Gradient Checkpointing |\n|------|-------|-------|-------|-------|-------|---------------------------|\n| Memory (GB)    | 18.82 | 18.85 | 18.94 | 19.08 | 22.77 | \\| 34.76                  |\n| Time (s) | 5.15  | 5.03  | 4.96  | 4.92  | 4.89  | \\| 4.83                   |\n\n**Table:** BP memory and time cost of Qwen 3-4B under sequence length 9k.\n\nFrom the above table, StreamBP introduces slight overhead compared to gradient checkpointing method, even though the total computational FLOPs is the same. The overhead is due to repetitive kernel launch and increased HBM throughput. As the partition size increases, the number of partitions decreases, leading to reduced BP time.\n\nWe hope that our response is satisfactory and addresses all your concerns. If there are any additional questions and concerns, please feel free to let us know during the author-reviewer discussion period. We are more than happy to provide further clarifications.\n\n**References**\n\n[1] Aashaka et al, \"Memory Optimization for Deep Networks\", ICLR 2021.\n\n[2] Liu et al, \"Blockwise Parallel Transformer for Large Context Models\", NeurIPS 2023."}, {"review_id": "0WoMLroMYN", "reviewer": "Reviewer_8uED", "summary": "Training large language models on long sequences demands substantial GPU memory during backpropagation. To address this, this paper introduces ​StreamBP, a memory-efficient and exact backpropagation method that processes sequences in chunks and sequentially accumulates gradients across them. By leveraging the causality of self-attention, StreamBP further reduces FLOPs during backpropagation by eliminating redundant attention computations. Compared to gradient checkpointing, StreamBP achieves faster speed and better memory efficiency during backpropagation.", "questions": "1. Could you elaborate the key difference between MsT and StreamBP?\n2. Line 50, why MsT does not apply to RL objectives?", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Strengths\n1. The paper is clearly written and well motivated. Although this is not my area of expertise, the diagrams help me quickly get up to understand the problem and follow the proposed solution.\n2. Compared to gradient checkpointing, StreamBP achieves faster speed and better memory efficiency.\n3. StreamBP is applicable to common scenarios such as SFT, GRPO, and DPO.\n\nWeaknesses\n1. The novelty is questionable, as the core idea of sequence partition already exists in MsT and Sequence parallel.\n2. Table 3 lacks the comparison between MsT and StreamBP, making it unclear whether StreamBP is more memory-efficient than StreamBP.", "was_revised": true, "final_justification": "Based on the author's response, I have re-evaluated the novelty of this paper and raised the rating to Borderline accept.", "rebuttal": "We thank the reviewer for the valuable comments. We now clarify the key differences between StreamBP and the baseline approaches.\n\n**Difference with MsT.** While MsT and StreamBP share certain similarities in terms of sequence partition, they exhibit fundamental differences that result in large efficiency gap and different application scope.\n\n1. Fundamental difference of the BP order that leads to large memory gap. MsT sequentially forwards all partitioned inputs *before* the full backward on the output (see their Algorithm 1). In comparison, StreamBP performs backward on the partitioned output *immediately* after each partitioned forward. This seemingly subtle difference is fundamental, as MsT still stores the full layer activation while StreamBP only stores the partitioned layer activation. This is because the activation value will only be released from memory after it has been utilized for gradient calculation. Therefore, the main memory saving of MsT comes from the reduced logits in the language modeling head layer, which is tailored to the SFT objective as we will analyze below. \n\n2. MsT is tailored to SFT; the technique cannot be applied to RLHF objectives, e.g., DPO. MsT’s forward (see their Algorithm 2) and backward (see their Algorithm 4) process of the language modeling head are specifically tailored to SFT objective. Specifically, MsT backwards the loss of one sequence partition at a time for avoiding storing the full logits. However, the technique is not applicable to RLHF objectives such as DPO, due to the linear loss decomposition idea of MsT (i.e., $L = L_1+\\ldots+L_D$) and the fact that the nonlinear sigmoid loss of DPO is not linearly decomposable across the sequence dimension. By contrast, StreamBP has a fundamentally different design idea, i.e., linear decomposition of the gradient computation; see our equation (1). This main design idea of StreamBP allows to derive specific forward and backward processes for RLHF objectives, which involves gradient analysis of nonlinear operator.\n\n3. StreamBP introduces nontrivial partitioning of attention computation, which is not present in MsT. This design in attention part improves both memory and computation efficiency of StreamBP. In particular, MsT adopts the standard attention implementation, which stores the full activations in attention layer including full $Q$, $K$, $V$ matrices. In comparison, StreamBP partitions the attention computation along the sequence dimension (see our Figure 2 and Equation 13) and reduces the activation values by approximately $D$ times, where $D$ is the number of partitions. This partition design is nontrivial due to the causal feature of the attention operator, where the output at position $t$ depends not only on the current position of input, but also on all the input at previous positions. Moreover, StreamBP reduces the computational FLOPs by leveraging the causal structure of language models. Namely, it avoids computing the key and value behind the current partition position. As shown in Table 1, the fine-grained attention computation allows StreamBP to achieve remarkable acceleration over MsT and even the vanilla gradient checkpointing method.\n\n4. StreamBP proposes a general principle for reducing the memory cost of training neural networks. Unlike MsT that uses loss decomposition (i.e., $L = L_1+\\ldots+L_D$), we would like to highlight that StreamBP’s main design idea of decomposing Jacobian-vector product (see our equation (1)) is applicable to general neural network training. For any transformation inside the model, if computing one partition of the output requires storing much less activation values than that for computing full output, then one can apply StreamBP to greatly reduce the memory cost. \n\n**Difference between StreamBP and sequence parallel.** Sequence parallel (SP) is specifically designed for distributed training with multiple available GPUs, while StreamBP is applicable for both single GPU and multi-GPUs settings. Specifically, SP partitions the activation across different GPUs, where each GPU receives one activation partition evaluated on the same input data. Consequently, SP requires a sufficient number of GPUs to scale to long sequences. Additionally, as different segments of the activation are stored in different GPUs, SP requires additional communication of activation values between multiple GPUs. By contrast, StreamBP sequentially processes all partitions on the same GPU, storing only one partition at a time. This design allows StreamBP to facilitate long sequence training even in *a single GPU* setting, while also allowing further scaling of sequence length when multiple GPUs are available. \n\nIn summary, StreamBP differs fundamentally from both MsT and sequence parallel in algorithmic design, memory efficiency, and applicable scenarios. These distinctions support our claim of StreamBP’s novelty.\n\n**\"Table 3 lacks the comparison between MsT and StreamBP\".** We present the memory comparison between StreamBP and MsT in Figure 4 rather than Table 3, as MsT is specifically designed for SFT as outlined above. When the sequence length scales up, there is a clear gap in memory consumption between StreamBP and MsT. The memory gap is caused by layer activation values, as we discussed above (Difference with MsT, 1&3).\n\n**\"MsT does not apply to RL objectives\".** Please see the discussion above (Difference with MsT, 2).\n\nWe hope that our response is satisfactory and addresses all your concerns. Should you have any further questions and concerns, please feel free to let us know during the author-reviewer discussion period. We are more than happy to provide further clarifications."}, {"review_id": "EBMAzUGXu2", "reviewer": "Reviewer_RVuF", "summary": "The paper discusses a method to improve the memory efficiency of LLMs post-training  by partitioning the input of each layer (MLP, Attention, Output layer)  into chunks and feeding them sequentially to the layer operation. This is done in the forward pass, but also during the backward pass (using checkpointing) to save memory on the intermediate activations. For the attention operation the model precomputes the keys and values, which are reused for all chunks of the queries. The method brings significant memory and time savings when training on long sequences compared to gradient checkpointing and the competitor method MsT. This enables for example to fit the forward and backward pass of Qwen3-14B with a sequence of 48k tokens in a single 80GB GPU, while MsT and checkpointing allow only for 9k and 1.7k tokens respectively (See Figure 3). The proposed method is also faster than MsT and checkpointing thanks to a more efficient attention operation using half the flops.", "questions": "I am willing to increase the score depending on the answers to my questions, Q1 in particular.\n\n**Questions and Comments**\n\nQ1 The authors mention a reduction to half the FLOPs for their attention computation. I feel this is a core part of the approach since this is where the time savings come from, and should be discussed more and maybe put also in the contributions. I also thought that using flash-attention would already result in such a reduction in FLOPS, is that not the case? I feel it would also be useful to state clearly which attention implementation is used in the text.\n\nQ2 The authors go into details into how the implementation of the method differs between SFT, DPO and GRPO. However, it seems that the method can be applied in the same fashion whenever the loss can be written as a sum, which is the case for GRPO and SFT. In the case of DPO instead the loss is still a nonlinear function of a scalar ($s$) which can be written as a sum, so the method is applied to $s$ and then adjusted by using the derivative of the loss w.r.t. $s$. Describing these three cases in detail seems unnecessary and the exposition could be greatly simplified by describing how the procedure can be applied whenever the function to be differentiated is a function of a scalar sum. The handling of attention is a bit more nuanced and its dedicated section could be expanded.\n\nQ2 Lines 58-59, “nontrivial and intricate …” This part is not so clear. I would appreciate it if the authors would be more specific on this point, for example mentioning the chunking that is done by the method? \n \nQ3 I found it interesting that the memory peaks are at the end of the second forward pass and not the first. I assume this is due to the fact that pytorch caches some memory to be reused during the forward pass. Maybe the authors can comment a bit on this point?\n\nMinor:\n\nM1 Figure 1: what is the batch size (I assumed one)? Maybe this could be indicated in the caption.\n \nM2 Typo above eq. (13): I think $S^{(i)}$ should be $M^{(i)}$.\n\nM3 Appendix A.1: $L$ is not defined.", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 2, "strengths_and_weaknesses": "**Strengths**\n\nS1 The method is conceptually simple and could be easily implemented without requiring custom GPU kernels. it is essentially a smarter checkpointing tailored to LLMs training, which takes advantage of the causal nature of the transformers architecture and divides the computation over the time dimension into chunks which are computed sequentially in order to save memory.\n\nS2 The method shows great advantage in memory and speed compared to checkpointing and a competition method named MsT (see end of my summary). It is quite remarkable that such a gain can be obtained without using custom CUDA or triton kernels, but just by cleverly chunking the input sequence.\n\nS3 The method is tailored for long sequence post-training scenarios (SFT, DPO and GRPO) that use long sequences, which is timely given the recent rise of reasoning models using chains of thought and thus generating very long sequences.\n\n**Weaknesses**\n\nW1 An experiment where the authors train a model with SFT, DPO, or GRPO is missing. Having this experiment would give additional strength to the proposed approach, since even though in infinite precision the derivative is computed in the same way as in standard back-propagation or checkpointing, numerical errors might make it more or less accurate to use the proposed approach. The experiment could also  show if this method enables post-training with long sequences on limited hardware and in a reasonable time.\n\nW2 Exposition can be improved. I feel that the exposition of the idea when applied to the last layer outputting the logits in the case of SFT, DPO and GRPO to be given too much spotlight, since the applications to those cases is relatively straightforward from the main idea. I would have liked more focus on the attention part which is less straightforward and it’s where the time saving comes (See also the questions section).\n\nW3 Limited novelty. The approach can be seen as a direct follow up of MsT [1] where the chunking of the input sequence is also proposed but is done at a lower granularity. This is a minor weakness though since the improvements are still remarkable even when compared to MsT.\n\n**References:**\n[1] Cheng Luo, Jiawei Zhao, Zhuoming Chen, Beidi Chen, and Anima Anandkumar. Mini-sequence transformer: Optimizing intermediate memory for long sequences training. arXiv preprint arXiv:2407.15892, 2024.", "was_revised": true, "final_justification": "The authors have adequately addressed all three points that I raised. In particular clarifying the point regarding the FLOPs saving compared to competitor methods and testing the robustness of the approach after several steps of fine-tuning. Therefore I changed the score from 4 to 5.", "rebuttal": "We thank the reviewer for the valuable comments. We address all your concerns below.\n\n**Reduction in FLOPS (Q1).** The FLOPs reduction is due to that StreamBP avoids computing the most of the upper triangular entries of the $QK^{\\top}$ matrix, which are not used in the attention computation due to the causal structure of language models. Specifically, the attention implementation is based on PyTorch's `scaled_dot_product_attention` with backend being flash attention. Based on the official document, the full $QK^{\\top}$ matrix will be computed as long as custom attention mask is required, which is mandatory when using techniques such as padding or sliding window attention. When using StreamBP, since the forward is partitioned along sequence dimension, we are able to perform fine-grained control of the attention computation. Specifically, for the sequence partition $(t_1, t_2)$, we compute `Q[t1:t2, :] K[:t2, :].transponse()` rather than `Q[t1:t2, :] K[:, :].transponse()` for avoiding unnecessary computations. The computation is based on flash attention without materializing the partitioned pre-attention score in the memory. We remark that StreamBP offers flexibility for fine-grained control of attention over sequence level without the need to implement a customized kernel. We will provide more discussions on the reduced FLOPs in the next version of our manuscript.\n\n**Numerical stability of StreamBP in model training on SFT/GRPO/DPO (W1).** Thank you for the suggestion. We have verified numerical accuracy of StreamBP in Table 6, where StreamBP obtains the gradient that has very close precision error as standard BP. In the following tables, we present the evaluation loss of training Qwen 3-4B model for 500 steps with SFT, DPO, and GRPO, respectively.\n\n| Step | Base   | StreamBP | Difference |\n|:----:|:------:|:--------:|:----------:|\n| 1    | 5.7821 | 5.7821   | +0.0000    |\n| 100  | 4.9235 | 4.9256   | +0.0021    |\n| 200  | 4.3128 | 4.3094   | -0.0034    |\n| 300  | 3.9896 | 3.9929   | +0.0033    |\n| 400  | 3.7914 | 3.7947   | +0.0033    |\n| 500  | 3.6670 | 3.6631   | -0.0039    |\n\n**Table:** Evaluation loss of SFT on dataset Capybara.\n\n| Step | Base Loss | StreamBP Loss | Difference |\n|:----:|:---------:|:-------------:|:----------:|\n| 1    | 1.7023    | 1.7023        | +0.0000    |\n| 100  | 1.0874    | 1.0895        | +0.0021    |\n| 200  | 0.9378    | 0.9346        | -0.0032    |\n| 300  | 0.8759    | 0.8792        | +0.0033    |\n| 400  | 0.8360    | 0.8399        | +0.0039    |\n| 500  | 0.8007    | 0.7963        | -0.0044    |\n\n**Table:** Evaluation loss of DPO on dataset Ultrafeedback.\n\n| Step | Base Loss | StreamBP Loss | Difference  |\n|:----:|:---------:|:-------------:|:----------------------------:|\n| 1 | 0.0127 | 0.0127 |  +0.0000  |\n| 100 | 0.0110 | 0.0108 |  -0.0002  |\n| 200 | 0.0048 | 0.0044 |  -0.0004  |\n| 300 | 0.0056 | 0.0056 |  0.0000  |\n| 400 | 0.0048 | 0.0049 |  +0.0001  |\n| 500 | 0.0043 | 0.0042 |  -0.0001  |\n\n**Table:** Evaluation loss of GRPO on dataset tldr.\n\nOne can see that the loss of StreamBP is very close to that of the standard BP. The tiny difference is due to different order of float operations; see Appendix B.1 for more details. We will add this additional verification experiment in the next version.  \n\n**Exposition can be improved (W2 & Q2).** Thank you for the suggestion. We plan to move the derivation of the GRPO loss to the appendix, as it becomes similar to that of the SFT loss once we observe that the GRPO loss is linearly decomposable along the token dimension. However, we would like to retain the derivation of the DPO loss in the main text, as it involves nontrivial observations of the gradient operator and requires a more careful decomposition. With the space saved, we intend to provide additional details on the attention part as suggested by the reviewer.\n\n**Limited novelty; can be seen as a direct follow up of MsT (W3).** Though it seems a little arguing as the reviewer mentioned this is a minor weakness, we would like to briefly clarify the fundamental differences betwee MsT and StreamBP. 1) MsT sequentially forwards all partitioned inputs before the full backward on the output, while StreamBP performs backward on the partitioned output immediately after each partitioned forward. This seemingly subtle difference results in large memory gap in terms of saving the layer activation. 2) Fundamentally different from the design idea of StreamBP (i.e., linear decomposition of the gradient computation; see our equation (1)), MsT's main design relies on a linear loss decomposition of different partitions (i.e., $L = L_1+\\ldots+L_D$), which is not applicable when the loss is nonlinear across sequence dimension, e.g., the DPO objective. 3) StreamBP involves nontrivial partitions of the attention computation, while this part is not present in MsT. In summary, though these two approaches share certain similarities in terms of sequence partition, we believe that StreamBP is fundamentally different from MsT. For more details about the difference, we refer to our response to Reviewer 8uED.\n\n**Lines 58-59, “nontrivial and intricate …” (another Q2).** We apologize for any possible confusion. We meant to express that though the main idea of linearly decomposing the gradient computation seems simple, applying it to the concrete transformer (last layer, MLP, and attention) requires nontrivial and intricate developments, especially in the attention computation. We will clarify this part in the next version. \n\n**Peak memory appears on  at the end of the second forward pass, not the first (Q3).** The additional memory cost at the end of the second forward pass is the model gradient, which is accumulated during the first backward pass (annotated in Figure 1). Please refer to Appendix C.1 for a detailed explanation of the memory profile.\n\n**Minor:** \n* **M1.** Yes, the batch size is 1. We will add the information in the caption.\n* **M2.** It should be $S^{(i)}$, which is the pre-attention score.\n* **M3.** $L$ can be arbitrary objective that depends on the output $Z$. We will provide clarification in the manuscript.\n\nWe hope that our response is satisfactory and addresses all your concerns. If there are any additional questions and concerns, please feel free to let us know during the author-reviewer discussion period. We are more than happy to provide further clarifications."}]

Author rebuttal:

Meta-review: On this paper, the authors propose a method to improve the memory efficiency of LLM training, by proposing an update to the backpropagation step. They do so, by first partitioning the input of each layer (e.g., Attenion, output layer, MLP) and feeding them sequentially both during the forward and the backward pass. The main idea is that then the entire values do not need to be recomputed during the gradient checkpointing. The authors show that the computed gradients are exact, and that the method works on different training strategies such as SFT, DPO and GRPO (RL). They also show that by using this method it is possible to go to quite extreme training tokens, in particular, they are able to train a Qwen3-14B with a sequence of 48k tokens in a single 80GB GPU

The paper initially got borderline reviews, 3 Borderline Accepts and  1 Borderline Reject. The reviewers praised the experimental section, its generalization to different tasks, good overall performance and the writing of the paper. However, they were divided into the novelty of the work. At the same time, the reviewers found some weaknesses in the paper when it comes to evaluating only in Qwen (the authors showed in rebuttal evaluation in Gemma), asking for some clarification which the authors properly addressed in the rebuttal, more comparisons with prior works (which the authors did in the rebuttal), clarification in number of flops and asking for experiments in numerical stability of the method in model training on SFT/GRPO/DPO. The authors gave good explanation and responses to all these points in the rebuttal.

After the discussion stage, all the reviewers increased their scores, with the final scores  being 3 Accepts and 1 Borderline Accept, unanimously recommending the paper to be accepted. I agree with the reviewers about the decision and justification, and urge the authors to implement all the changes in the camera-ready version of the paper. Congratulations to the authors for a very nice NeurIPS paper!
Decision: Accept (poster)