# ProTrain: Efficient LLM Training via Memory-Aware Techniques

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

## Key facts

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

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

- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [OPT: Open Pre-trained Transformer Language Models](https://mlsanity.com/paper/2205.01068/) (2022)
- [Mistral 7B](https://mlsanity.com/paper/2310.06825/) (2023)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://mlsanity.com/paper/1910.02054/) (2019)
- [Mixed Precision Training](https://mlsanity.com/paper/1710.03740/) (2017)
- [Training Deep Nets with Sublinear Memory Cost](https://mlsanity.com/paper/1604.06174/) (2016)
- [PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel](https://mlsanity.com/paper/2304.11277/) (2023)
- [ZeRO-Offload: Democratizing Billion-Scale Model Training](https://mlsanity.com/paper/2101.06840/) (2021)
- [Reducing Activation Recomputation in Large Transformer Models](https://mlsanity.com/paper/2205.05198/) (2022)
- [ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning](https://mlsanity.com/paper/2104.07857/) (2021)
- [FP8 Formats for Deep Learning](https://mlsanity.com/paper/2209.05433/) (2022)
- [Self-attention Does Not Need O(n2) Memory](https://mlsanity.com/paper/2112.05682/) (2021)
- [GSPMD: General and Scalable Parallelization for ML Computation Graphs](https://mlsanity.com/paper/2105.04663/) (2021)
- [Colossal-AI: A Unified Deep Learning System For Large-Scale Parallel Training](https://mlsanity.com/paper/2110.14883/) (2021)
- [Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization](https://mlsanity.com/paper/1910.02653/) (2019)
- [Galvatron: Efficient Transformer Training over Multiple GPUs Using Automatic Parallelism](https://mlsanity.com/paper/2211.13878/) (2022)
- [Optimal checkpointing for heterogeneous chains: how to train deep neural networks with limited memory](https://mlsanity.com/paper/1911.13214/) (2019)
- [PatrickStar: Parallel Training of Pre-trained Models via Chunk-based Memory Management](https://mlsanity.com/paper/2108.05818/) (2021)
- [TFLMS: Large Model Support in TensorFlow by Graph Rewriting](https://mlsanity.com/paper/1807.02037/) (2018)
- [Harmony: Overcoming the Hurdles of GPU Memory Capacity to Train Massive DNN Models on Commodity Servers](https://mlsanity.com/paper/2202.01306/) (2022)
- [Rockmate: an Efficient, Fast, Automatic and Generic Tool for Re-materialization in PyTorch](https://mlsanity.com/paper/2307.01236/) (2023)

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

- [Pie: Pooling CPU Memory for LLM Inference](https://mlsanity.com/paper/2411.09317/) (2024)
- [Gated Differentiable Working Memory for Long-Context Language Modeling](https://mlsanity.com/paper/2601.12906/) (2026)
- [A Survey of LLM × DATA](https://mlsanity.com/paper/2505.18458/) (2025)
- [EMA: Elastic and Performance Transparent Memory Across GPUs](https://mlsanity.com/paper/2609.27040/) (2026)

## Abstract (arXiv)

Memory pressure has emerged as a dominant constraint in scaling the training of large language models (LLMs), particularly in resource-constrained environments. While modern frameworks incorporate various memory-saving techniques, they often expose low-level configuration knobs that require manual tuning and specialized system expertise. This not only adds engineering overhead but also risks suboptimal hardware utilization when misconfigured. This paper introduces ProTrain, a novel training system that automatically tailors memory management policies to the model architecture and underlying hardware resources, eliminating the need for manual intervention. The core of ProTrain is its automated memory management that abstracts complex memory management strategies into a few tunable configuration parameters, allowing searches for optimal parameter settings using cost models. ProTrain is equipped with a runtime profiler that provides precise estimates of latency, memory usage, and I/O bandwidth to build high-fidelity cost models. ProTrain does not change the training algorithm and thus does not compromise accuracy. Experiments show that ProTrain improves training throughput by 1.43× to 2.71× compared to the state-of-the-art training systems.