Research map / Architectures and efficiency / RNNs and normalization
StreamBP: Memory-Efficient Exact Backpropagation for Long Sequence Training of LLMs
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes StreamBP, an exact backpropagation method that linearly decomposes the chain rule along the sequence dimension to cut activation and logit memory for long-sequence LLM training (SFT, GRPO, DPO).
Problem
Backpropagation memory for storing activations and logits grows with sequence length in LLM training on long reasoning traces, and gradient checkpointing still stores full layer activations and logits.
Approach
Partition the output along the sequence dimension and accumulate exact gradients chunk by chunk via a linear decomposition of the chain rule, exploiting causal attention so that each chunk only uses keys/values up to its position, reducing both memory and FLOPs compared with checkpointed BP.
Key results
- StreamBP scales maximum BP sequence length by 2.8-5.5x over gradient checkpointing on a single 80GB A800 for Qwen 3-8B/14B/32B.
- StreamBP achieves 5-5.6x larger maximum sequence length than checkpointing under DeepSpeed ZeRO-2 on 8 GPUs for Qwen 3-8B.
- For Qwen 3-4B, StreamBP BP time is 6-13% lower than checkpointing at sequence lengths of 12k and above.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace gradient checkpointing with StreamBP | Qwen 3 models trained with SFT, GRPO, or DPO · Single A800-80GB, configurations in maximum-sequence-length table | maximum sequence length: 2.8-5.5× | Helped |
| Use StreamBP instead of gradient checkpointing | Qwen 3-4B · BP time, sequence lengths 6k–27k | BP time: --4.4%; 0.4%; 6.0%; 7.4%; 10.5%; 10.9%; 12.9% | Mixed |
| Use StreamBP instead of gradient checkpointing | Qwen 3-8B · ZeRO-2, 6k–21k sequence lengths | per-sample BP time (seconds): 8.4 -> 6.4; 9.6 -> 8.3; 12.2 -> 10.8; 18.5 -> 14.8; 22.9 -> 17.2; 26.3 -> 20.7 | Helped |
| Increase batch size while using StreamBP | Qwen 3-4B · Sequence length 9000 | per-sample BP time: 7.03; 6.82; 6.63; 6.47 | Helped |
| Increase partition count (ablation) | StreamBP, linear example · Linear example | peak memory (GB): 36.01 -> 13.25; 12.47; 12.20; 12.07; 11.99 | Helped |
| Increase partition count (ablation) | StreamBP, linear example · Linear example | time cost (seconds): 14.48 -> 14.50; 14.70; 14.79; 15.45; 18.42 | Hurt |
| Use a partition size that is too small (ablation) | StreamBP transformer layers · Long-sequence training | BP time | Hurt |
| Use a relatively large partition size (ablation) | StreamBP transformer layers · Long-sequence training | training efficiency | Helped |
| Use StreamBP under BF16 instead of standard BP under BF16 | Qwen 3 language-model head and transformer layers · Gradient precision analysis | relative gradient error: 1.43% vs 1.47%; 6.63% vs 6.59% | No clear effect |
Surprising findings
- Gradient difference of StreamBP in bfloat16 matches standard BF16 BP error levels, showing no added precision loss.
Limitations
- Does not support MoE or multimodal models currently.
- Partition size introduces HBM overhead and kernel launch cost at small partition sizes.
- Distributed design adds memory for caching one transformer layer per GPU.
Experimental setup
Metrics: maximum sequence length peak memory backpropagation time gradient absolute and relative error
Baselines: gradient checkpointing baseline without checkpointing MsT
Builds on: gradient checkpointing flash attention DeepSpeed ZeRO GRPO DPO
Closest papers by meaning
- Mini-Sequence Transformers: Optimizing Intermediate Memory for Long Sequences Training (NeurIPS 2024 · similarity 0.88)
- Scaling Instruction-tuned LLMs to Million-token Contexts via Hierarchical Synthetic Data Generation (ICLR 2025 · similarity 0.88)
- xLSTM 7B: A Recurrent LLM for Fast and Efficient Inference (ICML 2025 · similarity 0.88)
- ShiftAddLLM: Accelerating Pretrained LLMs via Post-Training Multiplication-Less Reparameterization (NeurIPS 2024 · similarity 0.88)
- EE-LLM: Large-Scale Training and Inference of Early-Exit Large Language Models with 3D Parallelism (ICML 2024 · similarity 0.87)
- InfLLM: Training-Free Long-Context Extrapolation for LLMs with an Efficient Context Memory (NeurIPS 2024 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.