StreamBP: Memory-Efficient Exact Backpropagation for Long Sequence Training of LLMs

We thank the reviewer for the valuable comments. We address all your concerns below.

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.

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.

Table: Evaluation loss of SFT on dataset Capybara.

Table: Evaluation loss of DPO on dataset Ultrafeedback.

Table: Evaluation loss of GRPO on dataset tldr.

One 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.

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.

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.

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.

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.

Minor:

• M1. Yes, the batch size is 1. We will add the information in the caption.
• M2. It should be $S^{(i)}$, which is the pre-attention score.
• M3. $L$ can be arbitrary objective that depends on the output $Z$. We will provide clarification in the manuscript.

We 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.

We thank the reviewer for the valuable comments. We now clarify the key differences between StreamBP and the baseline approaches.

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.

1. 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.

2. 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.

3. 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.

4. 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.

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.

In 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.

"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).

"MsT does not apply to RL objectives". Please see the discussion above (Difference with MsT, 2).

We 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.

Thank you for recognizing the novelty of StreamBP and for increasing your score. In the next version of our manuscript, we will clarify the differences between StreamBP and MsT, particularly in terms of the backpropagation order, applicable training objectives, and attention computation partitioning. We will also elaborate on the distinctions between StreamBP and sequence parallelism, including their respective design choices and GPU requirements.

We thank the reviewer for the valuable comments. We now address your concerns and questions below.

Comparison with MoNET [1]

• 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.

• 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.

Comparison with Block Parallel Transformer (BPT) [2]

• 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.
• 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.
• StreamBP is compatible with most of the fused attention kernels, e.g., flash attention, and thereby has much less HBM throughput than BPT.

We 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.

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:

Table: BP memory and time cost of Qwen 3-4B under sequence length 9k.

From 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.

We 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.

References

[1] Aashaka et al, "Memory Optimization for Deep Networks", ICLR 2021.

[2] Liu et al, "Blockwise Parallel Transformer for Large Context Models", NeurIPS 2023.

We thank the reviewer for the valuable comments. We now address your concerns below.

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.

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.

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.

Table: BP memory cost (GB) of Llama 3.1-8B.

Table: BP memory cost (GB) of Gemma 3-1B.

Clearly, 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.

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.

We 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.