ProTrain: Efficient LLM Training via Automatic Memory Management
ML System;Memory Optimization;Data Parallelism;ZeRO;Gradient Checkpointing;Tensor Offloading
Yes
No ethics review needed.
1. Please provide some results about search overhead and the searched configuration.

2. Please discuss the novelty part. I think the technical contribution of the paper is somewhat restricted, as it relies on the application of ZeRO, Offloading and Activation Checkpointing. Their combination is not novel enough. Could you please highlight the challenge of combining them?
Thank you for submitting this paper to ICLR!

1. Important problem. Reducing memory consumption in LLM training is important and challenging.

2. Clear written, well-organized, and easy to follow.
This paper introduced ProTrain, a novel training system designed to simplify the training process through automatic memory management. ProTrain highlights the significance of precise memory usage and runtime data gathered through memory-aware, model-wise profiling to build high-fidelity cost models, along with the careful abstraction of configuration parameters from memory management strategies to automate optimal configuration search.The basic idea of ProTrain is to abstract memory management strategies into a few tunable configuration parameters. ProTrain then builds runtime and memory usage estimators that quantify the impacts of these configuration parameters on training performance. These cost models, informed with accurate profiling information on latency, memory, and I/O bandwidth, allow ProTrain to search for the optimal memory management strategy that minimizes runtime while ensuring the peak memory consumption meets the hardware constraints.
1. Limited technical innovation. The technical contribution of the paper is somewhat restricted, as it relies on the application of ZeRO, Offloading and Activation Checkpointing. Their combination is not novel enough, and the large search space is not well justified (can be easily pruned).

2. Small-scale experiments: The experiments are conducted on a single node server, with only 4 GPUs. Although authors have some discussion on the two V100 nodes, the experiments are not sufficient to support the LLM training scenario. The paper should provide more extensive experiments on larger scales to demonstrate the scalability of the proposed method. 

3. Unfair baseline comparison: The paper compares the proposed method with the vanilla FSDP (i.e., ZeRO-3 etc). However, Pytorch already support selective activation checkpointing. The paper should compare with the Pytorch's combination of FSDP, Offloading and Activation Checkpointing.
Yes
No ethics review needed.
1. for optimizer in the backward, curious how does it compose with gradient norm clipping? for LLM training, it seems quite critical to clip gradients base on total norms to avoid numeric issues. If I understand it correctly, we perform optimizer step for each gradient or layer now without knowing the total gradient norms

2. for "Some frameworks fail to scale model sizes with more GPUs, primarily due to inefficiencies in handling model initialization across devices.", curious what are the foundmental issue for FSDP to only support 1B model at most? Is it init model on gpu before applying fsdp?

3. it would be great to summarize the exact combination of different tricks when they reach the best memory and perf.
1. The search space covers wide spectrum of memory technics including ZeRO, param offloading, activation offloading, activation checkpoint, gpu persistent parameters, and optimizer in the backward

2. Runtime estimator and peak memory estimator are foudemental work that can be useful in general, including OOM prevention and offline simulation
This paper proposed a cost model to evaluate different memory technics to achive best perf under memory constraint. The memory technics include 1) parameters offloading, 2) parameters all-gather and gradients reduce-scatter, 3) activation checkpoint, 4) activation offloading. It also proposed a dual-chunk heuristics that persist first layer on gpu to avoid exposed all-gathers. 

The author also built 2 fundemental components to support cost model: 1) runtime estimator for forward, backward, gpu/cpu optimizer, 2) peak memory estimator
1. The explanation for maximun trainable model size seem to be vague. "Some frameworks fail to scale model sizes with more GPUs,
primarily due to inefficiencies in handling model initialization across devices.". It would be great to be concrete about this. How exactly each baseline is configured and why they fail. Is it because of meta/cpu/gpu init, or cpu offloading

2. For training throughput, it would be great to cover the configuration for baseline - whether they use cpu offloading, what are the activation checkpoints, how parameters are grouped for all-gather and reduce-scatter
Yes
No ethics review needed.
1. **Persistent Chunks?** Given that the model states for the initial chunk are only required at the start and end of the training pipeline, what is the advantage of keeping them persistently on the GPU? In two consecutive iterations, it’s understandable that updating these states on the GPU is efficient since they’re reused right after the first iteration. However, why not evict these states during the forward pass to free up space for other data, swapping them back only at the end of the backward pass to optimize memory usage, just like the other chunks?
2. **Difference from LRU?** Although the dual observations on managing model states and activations are valuable, it’s unclear how this approach offers distinct benefits over a computation-aware Least Recently Used (LRU) caching strategy, which also prioritizes recent data retention on the GPU within the stacked training pipeline. Could you clarify this difference?
3. **Portability Concerns** Since ProTrain is built independently on top of PyTorch, how can developers who are already using frameworks like DeepSpeed or Colossal-AI integrate ProTrain into their existing codebases?
1. The dual-chunk design for model states and interleaved organization for activation management effectively reduce memory overhead by transforming memory management into a configuration task with a few tunable parameters.
2. The authors implement a memory and computation profiling system that supports automated memory management during training, optimizing resource allocation without manual intervention.
3. ProTrain demonstrates promising results, not only in expanding the trainable model size but also in boosting throughput. Integrating this system into existing platforms could significantly enhance training efficiency.
This paper introduces **ProTrain**, a training framework aimed at optimizing large language model (LLM) training on memory-limited hardware through **automated** data management across GPUs and CPUs, using a few tunable parameters. ProTrain’s design is based on two key insights: (a) **Model states**—prioritizing GPU residence for the initial layers minimizes latency due to cold starts and backward computation order; (b) **Activations**—retaining the final layers on GPU facilitates immediate access during backward passes. For all other model states and activations, ProTrain employs **offloading (swapping)**  and **checkpointing** to reduce memory usage. ProTrain’s memory and computation profiling enables efficient overlap of memory operations with computation, maximizing hardware utilization. Evaluations show that ProTrain supports models up to **1.64x larger** and achieves a **1.43x throughput improvement** over state-of-the-art solutions.
1. The paper primarily addresses training with data parallelism; however, expanding ProTrain to pipeline parallelism would improve its applicability, especially as modern models rarely fit on a single GPU during training.
2. Forward pass efficiency also depends on batch size, and a limited parallel scale may not sufficiently overlap memory transfer with computation. A more detailed micro-benchmark could clarify how effectively ProTrain overlaps memory transfer with computation.
Yes
No ethics review needed.
1. What is the design overhead of ProTrain?

2. Why is managing activation swapping and gradient checkpointing operations at the transformer block level a better policy? Is it possible to achieve better thourghput with a more fine-grained management scheme, e.g., the smallest granularity of managing activation swapping and gradient checkpointing operations is 1/3 of the transformer block?
1. The paper works on an important topic.
2. The techniques are evaluated on various LLMs.
3. The techniques are tested on different hardware platforms.
Training large language models (LLMs) demands intricate memory management strategies. However, current memory management approaches often rely on manually tuned configurations, which results in suboptimal hardware utilization and performance. In this paper, the authors introduce ProTrain, an advanced training system that autonomously develops memory management policies to efficiently support the training of various LLMs across diverse hardware platforms. ProTrain reframes memory management as an optimization problem with a limited set of tunable parameters that are optimized using a cost model. Additionally, it incorporates a runtime profiler to provide precise estimations of latency, memory consumption, and I/O bandwidth, thus enhancing the accuracy of the cost model.
1. The techniques proposed by this paper cannot be applied on the real-world LLM training systems, since it includes off-loading. At most, these techniques can be used for only fine-tuning of the LLMs.

2. The paper lacks of motivaiton. Although this paper claims that "existing frameworks often depend on manual tuning of memory management settings, leading to inefficient hardware utilization and suboptimal performance". The paper did not show how bad existing frameworks are.

3. Why not LORA? Maybe LORA can achieve the same performance but consumes much less training overhead.

4. The paper is not well-written. Many key concepts are not well-introduced. For example, off-loading, and activation recomputation.
a high-throughput training framework for LLMs that intelligently balance memory usage and performance via automatic memory management
@inproceedings{
anonymous2024protrain,
title={ProTrain: Efficient {LLM} Training via Automatic Memory Management},
author={Anonymous},
booktitle={Submitted to The Thirteenth International Conference on Learning Representations},
year={2024},
url={https://openreview.net/forum?id=e1Z4NCQ146},
note={under review}
}
Training billion-scale large language models (LLMs) with just a few consumer-grade graphics cards is key to democratizing LLM access. However, existing frameworks often depend on manual tuning of memory management settings, leading to inefficient hardware utilization and suboptimal performance. 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. ProTrain features (1) automated memory management that abstracts complex memory management strategies into a few tunable configuration parameters and searches for optimal parameter settings using cost models and (2) 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$\times$ to 2.71$\times$ compared to the state-of-the-art training systems.
I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
ML System Memory Optimization Data Parallelism ZeRO Gradient Checkpointing Tensor Offloading
I certify that there is no acknowledgement section in this submission for double blind review.
/pdf/dd0498cafb3b10e5f436e08299ab01fe26174b2c.pdf
infrastructure, software libraries, hardware, systems, etc.
I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.
I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
/attachment/3ead9d02419d11fc5dcc4941b92f054a88bce152.pdf
ProTrain: Efficient LLM Training via Automatic Memory Management
ICLR 2025 Conference Submission
ICLR.cc/2025/Conference/Submission