# Cost-Efficient LLM Training with Lifetime-Aware Tensor Offloading via GPUDirect Storage

arXiv: https://arxiv.org/abs/2506.06472 · HTML: https://mlsanity.com/paper/2506.06472/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 24 of its 48 references that are in the map.
- First posted on arXiv 6 Jun 2025.

## Builds on (most cited first, 24 of 24)

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [LoRA: Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2106.09685/) (2021)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [Mixed Precision Training](https://mlsanity.com/paper/1710.03740/) (2017)
- [ZeRO-Offload: Democratizing Billion-Scale Model Training](https://mlsanity.com/paper/2101.06840/) (2021)
- [An Empirical Model of Large-Batch Training](https://mlsanity.com/paper/1812.06162/) (2018)
- [ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning](https://mlsanity.com/paper/2104.07857/) (2021)
- [Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization](https://mlsanity.com/paper/1910.02653/) (2019)
- [TorchTitan: One-stop PyTorch native solution for production ready LLM pre-training](https://mlsanity.com/paper/2410.06511/) (2024)
- [NetLLM: Adapting Large Language Models for Networking](https://mlsanity.com/paper/2402.02338/) (2024)
- [SuperNeurons: Dynamic GPU Memory Management for Training Deep Neural Networks](https://mlsanity.com/paper/1801.04380/) (2018)
- [Training Large Neural Networks with Constant Memory using a New Execution Algorithm](https://mlsanity.com/paper/2002.05645/) (2020)
- [Building AI Agents for Autonomous Clouds: Challenges and Design Principles](https://mlsanity.com/paper/2407.12165/) (2024)
- [Smart-Infinity: Fast Large Language Model Training using Near-Storage Processing on a Real System](https://mlsanity.com/paper/2403.06664/) (2024)
- [TFLMS: Large Model Support in TensorFlow by Graph Rewriting](https://mlsanity.com/paper/1807.02037/) (2018)
- [GPU-Initiated On-Demand High-Throughput Storage Access in the BaM System Architecture](https://mlsanity.com/paper/2203.04910/) (2022)
- [G10: Enabling An Efficient Unified GPU Memory and Storage Architecture with Smart Tensor Migrations](https://mlsanity.com/paper/2310.09443/) (2023)
- [Beyond the Memory Wall: A Case for Memory-centric HPC System for Deep Learning](https://mlsanity.com/paper/1902.06468/) (2019)
- [Basic Performance Measurements of the Intel Optane DC Persistent Memory Module](https://mlsanity.com/paper/1903.05714/) (2019)
- [Angel-PTM: A Scalable and Economical Large-scale Pre-training System in Tencent](https://mlsanity.com/paper/2303.02868/) (2023)
- [ZnG: Architecting GPU Multi-Processors with New Flash for Scalable Data Analysis](https://mlsanity.com/paper/2006.08975/) (2020)

## Built on by (most cited first, 3 of 3)

- [LazyTrain: Limited-resource Allocation toward Zero-waste Yield Optimization in Large Language Model Training](https://mlsanity.com/paper/2608.11919/) (2026)
- [GreedySnake: Accelerating SSD-Offloaded LLM Training with Efficient Scheduling and Optimizer Step Overlapping](https://mlsanity.com/paper/2512.17570/) (2025)
- [Horizon-LM: A RAM-Centric Architecture for LLM Training](https://mlsanity.com/paper/2602.04816/) (2026)

## Abstract (arXiv)

We present the design and implementation of a new lifetime-aware tensor offloading framework for GPU memory expansion using low-cost PCIe-based solid-state drives (SSDs). Our framework, TERAIO, is developed explicitly for large language model (LLM) training with multiple GPUs and multiple SSDs. Its design is driven by our observation that the active tensors take only a small fraction (1.7% on average) of allocated GPU memory in each LLM training iteration, the inactive tensors are usually large and will not be used for a long period of time, creating ample opportunities for offloading/prefetching tensors to/from slow SSDs without stalling the GPU training process. TERAIO accurately estimates the lifetime (active period of time in GPU memory) of each tensor with the profiling of the first few iterations in the training process. With the tensor lifetime analysis, TERAIO will generate an optimized tensor offloading/prefetching plan and integrate it into the compiled LLM program via PyTorch. TERAIO has a runtime tensor migration engine to execute the offloading/prefetching plan via GPUDirect storage, which allows direct tensor migration between GPUs and SSDs for alleviating the CPU bottleneck and maximizing the SSD bandwidth utilization. In comparison with state-of-the-art studies such as ZeRO-Offload and ZeRO-Infinity, we show that TERAIO improves the training performance of various LLMs by 1.47x on average, and achieves 80.7% of the ideal performance assuming unlimited GPU memory.