# 2BP: 2-Stage Backpropagation

arXiv: https://arxiv.org/abs/2405.18047 · HTML: https://mlsanity.com/paper/2405.18047/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 22 of its 36 references that are in the map.
- First posted on arXiv 28 May 2024.

## Builds on (most cited first, 22 of 22)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [PaLM: Scaling Language Modeling with Pathways](https://mlsanity.com/paper/2204.02311/) (2022)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [GLU Variants Improve Transformer](https://mlsanity.com/paper/2002.05202/) (2020)
- [Root Mean Square Layer Normalization](https://mlsanity.com/paper/1910.07467/) (2019)
- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://mlsanity.com/paper/1910.02054/) (2019)
- [ZeRO-Offload: Democratizing Billion-Scale Model Training](https://mlsanity.com/paper/2101.06840/) (2021)
- [ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning](https://mlsanity.com/paper/2104.07857/) (2021)
- [PipeDream: Fast and Efficient Pipeline Parallel DNN Training](https://mlsanity.com/paper/1806.03377/) (2018)
- [DAPPLE: A Pipelined Data Parallel Approach for Training Large Models](https://mlsanity.com/paper/2007.01045/) (2020)
- [Colossal-AI: A Unified Deep Learning System For Large-Scale Parallel Training](https://mlsanity.com/paper/2110.14883/) (2021)
- [Chimera: Efficiently Training Large-Scale Neural Networks with Bidirectional Pipelines](https://mlsanity.com/paper/2107.06925/) (2021)
- [PipeMare: Asynchronous Pipeline Parallel DNN Training](https://mlsanity.com/paper/1910.05124/) (2019)
- [Zero Bubble Pipeline Parallelism](https://mlsanity.com/paper/2401.10241/) (2023)
- [Hanayo: Harnessing Wave-like Pipeline Parallelism for Enhanced Large Model Training Efficiency](https://mlsanity.com/paper/2308.15762/) (2023)
- [PatrickStar: Parallel Training of Pre-trained Models via Chunk-based Memory Management](https://mlsanity.com/paper/2108.05818/) (2021)
- [Elixir: Train a Large Language Model on a Small GPU Cluster](https://mlsanity.com/paper/2212.05339/) (2022)

## Abstract (arXiv)

As Deep Neural Networks (DNNs) grow in size and complexity, they often exceed the memory capacity of a single accelerator, necessitating the sharding of model parameters across multiple accelerators. Pipeline parallelism is a commonly used sharding strategy for training large DNNs. However, current implementations of pipeline parallelism are being unintentionally bottlenecked by the automatic differentiation tools provided by ML frameworks. This paper introduces 2-stage backpropagation (2BP). By splitting the backward propagation step into two separate stages, we can reduce idle compute time. We tested 2BP on various model architectures and pipelining schedules, achieving increases in throughput in all cases. Using 2BP, we were able to achieve a 1.70x increase in throughput compared to traditional methods when training a LLaMa-like transformer with 7 billion parameters across 4 GPUs.