ProTrain: Efficient LLM Training via Automatic Memory Management

Authors: Hanmei Yang, Jin Zhou, Yao Fu, Xiaoqun Wang, Ramine Roane, Hui Guan, Tongping Liu
Primary area: infrastructure, software libraries, hardware, systems, etc.
Keywords: ML System, Memory Optimization, Data Parallelism, ZeRO, Gradient Checkpointing, Tensor Offloading

Abstract:
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.

Decision: reject
Reviewer scores: 5, 5, 6, 5

Reviews and discussion:
<thread>
creation date:
1730140497399

thread_type:
official_review

<thread_object>
creator: reviewer_eecs

<content>
summary:
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.

strengths:
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.

weaknesses:
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.

questions:
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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730484218988

thread_type:
official_review

<thread_object>
creator: reviewer_9uni

<content>
summary:
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.

strengths:
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.

weaknesses:
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.

questions:
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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730612172954

thread_type:
official_review

<thread_object>
creator: reviewer_sdyt

<content>
summary:
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

strengths:
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

weaknesses:
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

questions:
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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730661394489

thread_type:
official_review

<thread_object>
creator: reviewer_hsx9

<content>
summary:
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.

strengths:
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.

weaknesses:
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.

questions:
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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731896007894

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
## Common Questions

### Question-1: Novelty of ProTrain

Our novelty includes the following aspects: 

1. **Automation via Cost Models**: ProTrain automatically and efficiently identifies the best memory management policy for a given LLM architecture and hardware setup. Unlike existing systems, which require extensive manual tuning and grid search (often taking hours or days) to meet memory and performance needs, ProTrain completes this search within seconds (as shown in Common Question-2), streamlining the setup process and maximizing training throughput. Such automation is enabled by **two cost models** that estimate the memory consumption and runtime depending on configurations. 

2. **Novel Design Space**: ProTrain identifies a few new optimization opportunities when combining ZeRO, Offloading, and Activation Checkpointing in memory management and exposes them as additional tuning knobs to identify a better policy. (a) **ProTrain introduces a novel dual-chunk system for model states**. The design improves dynamic memory management by introducing persistent chunks and reusable chunk buffers. It also helps establish more accurate estimates on runtime and memory usage during profiling. (b) **ProTrain interleaves Swapping and Checkpointing**. The design is based on the rationale that, if the swapping overhead of a transformer block can be hidden by the forward and backward computation, the block would prefer swapping rather than checkpointing. Otherwise, the block should use checkpointing to reduce CPU-GPU bandwidth contention. 

3. **Model-wise Profiling**: ProTrain designed a novel model-wise runtime profiler that has two distinctive features. First, it is the first to consider temporary tensors and unhookable operators which can contribute up to 17.2% of peak memory usage but are often ignored by existing methods. Second, it supports profiling for large models that cannot fully fit into GPU memory by separating the memory consumption from static tensors, which are known as apriori, from the peak memory consumption estimation. The runtime profiler allows ProTrain to get precise peak memory consumption estimation for cost models, as demonstrated in Figure 5(b) and Figure 8 in Appendix D.4. 

### Question-2: Search overhead and searched configuration

Most of the time in obtaining the optimal memory management plan is spent on profiling, as **search itself takes only 0.06s on average**. Profiling time depends on the model's execution time; for instance, profiling Mistral-7B with batch size 4 takes 3.09s, while GPT-20B with batch size 4 takes 5.38s. These results were obtained using RTX 3090 GPUs.

The table below summarizes the automatically searched configurations with the best performance. It shows that model type, batch size, and hardware combinations can lead to varying configurations: 

| ID  | Model and Hardware                  | # of checkpointing blocks / # of blocks | # of swapping blocks | # of persistent chunks / # of chunks | # of chunk buffers |
| --- | ----------------------------------- | --------------------------------------- | -------------------- | ------------------------------------ | ------------------ |
| A   | GPT-1B BS=8 on 4 RTX 3090s or A100s | 0 / 32                                  | 0                    | 12 / 12                              | 0                  |
| B   | GPT-1B BS=64 on 4 RTX 3090s         | 24 / 32                                 | 2                    | 2 / 12                               | 3                  |
| C   | GPT-1B BS=64 on 4 A100s             | 0 / 32                                  | 0                    | 12 / 12                              | 0                  |
| D   | GPT-10B BS=8 on 4 RTX 3090s         | 48 / 48                                 | 0                    | 3 / 49                               | 46                 |
| E   | GPT-10B BS=8 on 4 A100s             | 0 / 48                                  | 0                    | 15 / 49                              | 3                  |

(to be continued...)
</content>
</thread_object>

</thread>
<thread>
creation date:
1732672356792

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
Thank you for taking the time to review our paper and provide valuable feedback. We have carefully addressed all the raised questions and provided detailed responses. As the author-reviewer discussion period is nearing its end, we kindly request any feedback you may have on our responses to ensure we have thoroughly addressed your concerns and further improved the work.

We hope that you will support our paper’s acceptance given the recognized strengths of the paper. Specifically, our work aims to make LLM training more accessible on resource-constrained GPUs, which we believe is a crucial advancement for the field. We have endeavored to address key challenges and have detailed these efforts in our submission. Thank you for your consideration and support!
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper proposes a novel system for automatically tuning memory management when training LLMs. While the paper tackles an important problem, the scope of the technical contributions is limited; furthermore, there are some concerns about the scope of the experiments and the baselines considered. The authors are encouraged to improve their paper either by significantly strengthening their experiments, or by expanding on the technical contribution to achieve large performance gains (e.g., expanding the search space of optimizations considered).