ProTrain: Efficient LLM Training via Memory-Aware Techniques

TL;DR

ProTrain is a system that automatically optimizes memory management for large language model training, improving throughput without manual tuning or accuracy loss.

Contribution

It introduces automated memory management policies tailored to hardware and model architecture, reducing engineering overhead and enhancing training efficiency.

Findings

ProTrain achieves 1.43× to 2.71× faster training throughput.

It uses cost models and runtime profiling for optimal memory policy search.

ProTrain maintains training accuracy while improving efficiency.

Abstract

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$\times$ to 2.71$\times$ compared to the state-of-the-art training systems.