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

• URL: http://arxiv.org/abs/2506.03077v1
• Date: Tue, 03 Jun 2025 16:54:15 GMT
• Title: StreamBP: Memory-Efficient Exact Backpropagation for Long Sequence Training of LLMs
• Authors: Qijun Luo, Mengqi Li, Lei Zhao, Xiao Li,
• Abstract summary: We propose a memory-efficient Backpropagation (BP) method called StreamBP.<n>StreamBP performs a linear decomposition of the chain rule along the sequence dimension in a layer-wise manner.<n>Compared to gradient checkpointing, StreamBP scales up the maximum sequence length of BP by 2.8-5.5 times larger.
• Score: 8.960494482210919
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• HelixPipe: Efficient Distributed Training of Long Sequence Transformers with Attention Parallel Pipeline Parallelism [14.067070576474086]
  As transformer sequence lengths grow, existing pipeline parallelisms incur suboptimal performance due to the quadratic attention computation and the substantial memory overhead.<n>We propose HelixPipe, a novel pipeline parallelism for long sequence transformer training.<n>It introduces attention parallel partition, which schedules attention computations of different micro batches across different pipeline stages in parallel, reducing pipeline bubbles.<n>It employs a two-fold first-in-last-out micro batch schedule to balance memory usage and overlap communication with fragmentation.
  arXiv  Detail & Related papers  (2025-07-01T03:11:18Z)
• TAH-QUANT: Effective Activation Quantization in Pipeline Parallelism over Slow Network [21.231881562816373]
  We introduce TAH-Quant (Tile-wise Adaptive Hadamard Quantization), a novel activation quantization framework designed specifically for pipeline parallelism.<n>Our approach integrates fine-grained tile-wise quantization for precise control, entropy-guided token-level adaptive bit allocation for optimal bit usage, and a Hadamard-based transform with pivot element swapping to effectively suppress quantization outliers.
  arXiv  Detail & Related papers  (2025-06-02T06:13:41Z)
• Scaling Recurrent Neural Networks to a Billion Parameters with Zero-Order Optimization [0.0]
  RNNs scale constant in FLOPs and GPU memory with increasing context length.<n>Transformers scale linearly in FLOPs and, at best, linearly in memory during generation.<n>Training large RNNs on long contexts remains impractical because standard optimization methods depend on Backpropagation Through Time (BPTT)
  arXiv  Detail & Related papers  (2025-05-23T13:04:06Z)
• LESA: Learnable LLM Layer Scaling-Up [57.0510934286449]
  Training Large Language Models (LLMs) from scratch requires immense computational resources, making it prohibitively expensive.<n>Model scaling-up offers a promising solution by leveraging the parameters of smaller models to create larger ones.<n>We propose textbfLESA, a novel learnable method for depth scaling-up.
  arXiv  Detail & Related papers  (2025-02-19T14:58:48Z)
• LoRC: Low-Rank Compression for LLMs KV Cache with a Progressive Compression Strategy [59.1298692559785]
  Key-Value ( KV) cache is crucial component in serving transformer-based autoregressive large language models (LLMs) Existing approaches to mitigate this issue include: (1) efficient attention variants integrated in upcycling stages; (2) KV cache compression at test time; and (3) KV cache compression at test time. We propose a low-rank approximation of KV weight matrices, allowing plug-in integration with existing transformer-based LLMs without model retraining. Our method is designed to function without model tuning in upcycling stages or task-specific profiling in test stages.
  arXiv  Detail & Related papers  (2024-10-04T03:10:53Z)
• SHERL: Synthesizing High Accuracy and Efficient Memory for Resource-Limited Transfer Learning [63.93193829913252]
  We propose an innovative METL strategy called SHERL for resource-limited scenarios. In the early route, intermediate outputs are consolidated via an anti-redundancy operation. In the late route, utilizing minimal late pre-trained layers could alleviate the peak demand on memory overhead.
  arXiv  Detail & Related papers  (2024-07-10T10:22:35Z)
• DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models [34.74093040678323]
  We introduce DeepSpeed-Ulysses, a novel, portable and effective methodology for enabling highly efficient and scalable LLM training. DeepSpeed-Ulysses at its core partitions input data along the sequence dimension and employs an efficient all-to-all collective communication for attention. Experiments show that DeepSpeed-Ulysses trains 2.5x faster with 4x longer sequence length than the existing method SOTA baseline.
  arXiv  Detail & Related papers  (2023-09-25T20:15:57Z)
• Blockwise Parallel Transformer for Large Context Models [70.97386897478238]
  Blockwise Parallel Transformer (BPT) is a blockwise computation of self-attention and feedforward network fusion to minimize memory costs. By processing longer input sequences while maintaining memory efficiency, BPT enables training sequences 32 times longer than vanilla Transformers and up to 4 times longer than previous memory-efficient methods.
  arXiv  Detail & Related papers  (2023-05-30T19:25:51Z)
• Winner-Take-All Column Row Sampling for Memory Efficient Adaptation of Language Model [89.8764435351222]
  We propose a new family of unbiased estimators called WTA-CRS, for matrix production with reduced variance. Our work provides both theoretical and experimental evidence that, in the context of tuning transformers, our proposed estimators exhibit lower variance compared to existing ones.
  arXiv  Detail & Related papers  (2023-05-24T15:52:08Z)
• Sequence Modeling with Multiresolution Convolutional Memory [27.218134279968062]
  We build a new building block for sequence modeling called a MultiresLayer. The key component of our model is the multiresolution convolution, capturing multiscale trends in the input sequence. Our model yields state-of-the-art performance on a number of sequence classification and autoregressive density estimation tasks.
  arXiv  Detail & Related papers  (2023-05-02T17:50:54Z)
• Mesa: A Memory-saving Training Framework for Transformers [58.78933015299703]
  We present Mesa, a memory-saving training framework for Transformers. Mesa uses exact activations during forward pass while storing a low-precision version of activations to reduce memory consumption during training. Experiments on ImageNet, CIFAR-100 and ADE20K demonstrate that Mesa can reduce half of the memory footprints during training.
  arXiv  Detail & Related papers  (2021-11-22T11:23:01Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.