# PETRA: Parallel End-to-end Training with Reversible Architectures

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

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Builds on 34 of its 47 references that are in the map.
- First posted on arXiv 4 Jun 2024.

## Builds on (most cited first, 30 of 34)

- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [Reformer: The Efficient Transformer](https://mlsanity.com/paper/2001.04451/) (2020)
- [Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour](https://mlsanity.com/paper/1706.02677/) (2017)
- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://mlsanity.com/paper/1910.02054/) (2019)
- [NICE: Non-linear Independent Components Estimation](https://mlsanity.com/paper/1410.8516/) (2014)
- [Training Deep Nets with Sublinear Memory Cost](https://mlsanity.com/paper/1604.06174/) (2016)
- [Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model](https://mlsanity.com/paper/2201.11990/) (2022)
- [A Downsampled Variant of ImageNet as an Alternative to the CIFAR datasets](https://mlsanity.com/paper/1707.08819/) (2017)
- [ResNet strikes back: An improved training procedure in timm](https://mlsanity.com/paper/2110.00476/) (2021)
- [ZeRO-Offload: Democratizing Billion-Scale Model Training](https://mlsanity.com/paper/2101.06840/) (2021)
- [The Reversible Residual Network: Backpropagation Without Storing Activations](https://mlsanity.com/paper/1707.04585/) (2017)
- [Decoupled Neural Interfaces using Synthetic Gradients](https://mlsanity.com/paper/1608.05343/) (2016)
- [i-RevNet: Deep Invertible Networks](https://mlsanity.com/paper/1802.07088/) (2018)
- [Training Neural Networks with Local Error Signals](https://mlsanity.com/paper/1901.06656/) (2019)
- [Memory-Efficient Pipeline-Parallel DNN Training](https://mlsanity.com/paper/2006.09503/) (2020)
- [Revisiting Neural Scaling Laws in Language and Vision](https://mlsanity.com/paper/2209.06640/) (2022)
- [PipeDream: Fast and Efficient Pipeline Parallel DNN Training](https://mlsanity.com/paper/1806.03377/) (2018)
- [Greedy Layerwise Learning Can Scale to ImageNet](https://mlsanity.com/paper/1812.11446/) (2018)
- [DAPPLE: A Pipelined Data Parallel Approach for Training Large Models](https://mlsanity.com/paper/2007.01045/) (2020)
- [Decoupled Greedy Learning of CNNs](https://mlsanity.com/paper/1901.08164/) (2019)
- [Revisiting Locally Supervised Learning: an Alternative to End-to-end Training](https://mlsanity.com/paper/2101.10832/) (2021)
- [PipeMare: Asynchronous Pipeline Parallel DNN Training](https://mlsanity.com/paper/1910.05124/) (2019)
- [Decoupled Parallel Backpropagation with Convergence Guarantee](https://mlsanity.com/paper/1804.10574/) (2018)
- [Reversible Vision Transformers](https://mlsanity.com/paper/2302.04869/) (2023)
- [torchgpipe: On-the-fly Pipeline Parallelism for Training Giant Models](https://mlsanity.com/paper/2004.09910/) (2020)
- [Can Forward Gradient Match Backpropagation?](https://mlsanity.com/paper/2306.06968/) (2023)
- [Training Neural Networks Using Features Replay](https://mlsanity.com/paper/1807.04511/) (2018)
- [On the Acceleration of Deep Learning Model Parallelism with Staleness](https://mlsanity.com/paper/1909.02625/) (2019)
- [Fully Decoupled Neural Network Learning Using Delayed Gradients](https://mlsanity.com/paper/1906.09108/) (2019)
- [Pipelined Backpropagation at Scale: Training Large Models without Batches](https://mlsanity.com/paper/2003.11666/) (2020)

## Built on by (most cited first, 6 of 6)

- [PosterGen: Aesthetic-Aware Multi-Modal Paper-to-Poster Generation via Multi-Agent LLMs](https://mlsanity.com/paper/2508.17188/) (2025)
- [Stabilizing Native Low-Rank LLM Pretraining](https://mlsanity.com/paper/2602.12429/) (2026)
- [Model Parallelism With Subnetwork Data Parallelism](https://mlsanity.com/paper/2507.09029/) (2025)
- [Learning without Global Backpropagation via Synergistic Information Distillation](https://mlsanity.com/paper/2510.03273/) (2025)
- [Learned Subspace Compression for Communication-Efficient Pipeline Parallelism](https://mlsanity.com/paper/2606.05484/) (2026)
- [Unbiased Approximate Vector-Jacobian Products for Efficient Backpropagation](https://mlsanity.com/paper/2602.14701/) (2026)

## Abstract (arXiv)

Reversible architectures have been shown to be capable of performing on par with their non-reversible architectures, being applied in deep learning for memory savings and generative modeling. In this work, we show how reversible architectures can solve challenges in parallelizing deep model training. We introduce PETRA, a novel alternative to backpropagation for parallelizing gradient computations. PETRA facilitates effective model parallelism by enabling stages (i.e., a set of layers) to compute independently on different devices, while only needing to communicate activations and gradients between each other. By decoupling the forward and backward passes and keeping a single updated version of the parameters, the need for weight stashing is also removed. We develop a custom autograd-like training framework for PETRA, and we demonstrate its effectiveness on CIFAR-10, ImageNet32, and ImageNet, achieving competitive accuracies comparable to backpropagation using ResNet-18, ResNet-34, and ResNet-50 models.