Cost-Efficient LLM Training with Lifetime-Aware Tensor Offloading via GPUDirect Storage
Ziqi Yuan, Haoyang Zhang, Yirui Zhou, Apoorve Mohan, I-Hsin Chung, Seetharami Seelam, Jian Huang
We present the design and implementation of a new lifetime-aware tensor offloadingframework for GPU memory expansion using low-cost PCIe-based solid-statedrives (SSDs). Our framework, TERAIO, is developed explicitly for large languagemodel (LLM) training with multiple GPUs and multiple SSDs. Its design is drivenby our observation that the active tensors take only a small fraction (1.7% onaverage) of allocated GPU memory in each LLM training iteration, the inactivetensors are usually large and will not be used for a long period of time, creatingample opportunities for offloading/prefetching tensors to/from slow SSDs withoutstalling the GPU training process. TERAIO accurately estimates the lifetime (activeperiod of time in GPU memory) of each tensor with the profiling of the first fewiterations in the training process. With the tensor lifetime analysis, TERAIO willgenerate an optimized tensor offloading/prefetching plan and integrate it into thecompiled LLM program via PyTorch. TERAIO has a runtime tensor migrationengine to execute the offloading/prefetching plan via GPUDirect storage, whichallows direct tensor migration between GPUs and SSDs for alleviating the CPUbottleneck and maximizing the SSD bandwidth utilization. In comparison withstate-of-the-art studies such as ZeRO-Offload and ZeRO-Infinity, we show thatTERAIO improves the training performance of various LLMs by 1.47× on average,and achieves 80.7% of the ideal performance assuming unlimited GPU memory.