ProTrain: Efficient LLM Training via Memory-Aware Techniques

训练大型语言模型需要极高的内存消耗。为了解决这个问题，现有的工作利用CPU和GPU的组合进行训练，例如ZeRO-Offload。这种技术在很大程度上民主化了十亿级模型的训练，使得只需要少量的消费级显卡就能进行训练。然而，根据我们的观察，现有的框架通常提供粗粒度的内存管理，并需要经验丰富的专家进行配置调优，导致硬件利用率和性能不够优化。本文提出了一种新的训练系统ProTrain，通过协调内存、计算和IO，智能平衡内存使用和性能。ProTrain通过基于块的模型状态管理和块状激活管理实现自适应内存管理，由Memory-Aware Runtime Profiler指导，无需用户干预。ProTrain不改变训练算法，因此不会影响准确性。实验表明，与SOTA训练系统相比，ProTrain提高了1.43倍至2.71倍的训练吞吐量。

ProTrain论文旨在解决大规模语言模型训练过程中的内存占用和性能问题，提出一种智能平衡内存使用和性能的训练系统。

ProTrain通过块状模型状态管理和块状激活管理实现自适应内存管理，同时利用内存感知的运行时分析器来协调内存、计算和IO，提高训练吞吐量。

ProTrain不需要用户干预，不会改变训练算法，不会影响模型精度。实验结果表明，ProTrain相比现有的最先进的训练系统，可以将训练吞吐量提高1.43倍至2.71倍。

近期相关研究包括ZeRO-Offload等利用CPU和GPU协同训练的技术，但现有框架的内存管理较为粗糙，需要经验丰富的专家进行配置调整，ProTrain提出的自适应内存管理方法是一种新的解决方案。