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

• URL: http://arxiv.org/abs/2506.06472v1
• Date: Fri, 06 Jun 2025 18:57:20 GMT
• Title: Cost-Efficient LLM Training with Lifetime-Aware Tensor Offloading via GPUDirect Storage
• Authors: Ziqi Yuan, Haoyang Zhang, Yirui Eric Zhou, Apoorve Mohan, I-Hsin Chung, Seetharami Seelam, Jian Huang,
• Abstract summary: TERAIO is a framework for GPU memory expansion using low-cost PCIe-based solid-state drives (SSDs)<n>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 large language iteration training process.<n>We show that TERAIO improves the training performance of various LLMs by 1.47x on average, and achieves 80.7% of the ideal performance.
• Score: 9.106167012987747
• License: http://creativecommons.org/licenses/by-nc-sa/4.0/
• Abstract: 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.

Related papers

• Minute-Long Videos with Dual Parallelisms [57.22737565366549]
  Diffusion Transformer (DiT)-based video diffusion models generate high-quality videos at scale but incur prohibitive processing latency and memory costs for long videos.<n>We propose a novel distributed inference strategy, termed DualParal.<n>Instead of generating an entire video on a single GPU, we parallelize both temporal frames and model layers across GPUs.
  arXiv  Detail & Related papers  (2025-05-27T11:55:22Z)
• Deep Optimizer States: Towards Scalable Training of Transformer Models Using Interleaved Offloading [2.8231000588510757]
  Transformers and large language models(LLMs) have seen rapid adoption in all domains. Training of transformers is very expensive and often hits a memory wall'' We propose a novel technique to split the LLM into subgroups, whose update phase is scheduled on either the CPU or the GPU.
  arXiv  Detail & Related papers  (2024-10-26T00:43:59Z)
• SSDTrain: An Activation Offloading Framework to SSDs for Faster Large Language Model Training [13.283682311968752]
  SSDTrain is an adaptive activation framework offloading to high-capacity GPU memory.<n>It is compatible with popular deep learning frameworks like PyTorch, Megatron, and DeepSpeed.<n>Results demonstrate that SSDTrain reduces 47% of the activation peak memory usage.
  arXiv  Detail & Related papers  (2024-08-19T14:09:48Z)
• vTensor: Flexible Virtual Tensor Management for Efficient LLM Serving [53.972175896814505]
  Large Language Models (LLMs) are widely used across various domains, processing millions of daily requests. Large Language Models (LLMs) are widely used across various domains, processing millions of daily requests.
  arXiv  Detail & Related papers  (2024-07-22T14:37:58Z)
• Hardware-Aware Parallel Prompt Decoding for Memory-Efficient Acceleration of LLM Inference [19.167604927651073]
  Auto-regressive decoding of Large Language Models (LLMs) results in significant overheads in their hardware performance. We propose a novel parallel prompt decoding that requires only $0.0002$% trainable parameters, enabling efficient training on a single A100-40GB GPU in just 16 hours. Our approach demonstrates up to 2.49$times$ speedup and maintains a minimal memory overhead of just $0.0004$%.
  arXiv  Detail & Related papers  (2024-05-28T22:19:30Z)
• FusionAI: Decentralized Training and Deploying LLMs with Massive Consumer-Level GPUs [57.12856172329322]
  We envision a decentralized system unlocking the potential vast untapped consumer-level GPU. This system faces critical challenges, including limited CPU and GPU memory, low network bandwidth, the variability of peer and device heterogeneity.
  arXiv  Detail & Related papers  (2023-09-03T13:27:56Z)
• FlexGen: High-Throughput Generative Inference of Large Language Models with a Single GPU [89.2451963569343]
  FlexGen is a generation engine for running large language model (LLM) inference on a single commodity GPU. When running OPT-175B on a single 16GB GPU, FlexGen achieves significantly higher throughput compared to state-of-the-art offloading systems. On the HELM benchmark, FlexGen can benchmark a 30B model with a 16GB GPU on 7 representative sub-scenarios in 21 hours.
  arXiv  Detail & Related papers  (2023-03-13T05:19:28Z)
• A Frequency-aware Software Cache for Large Recommendation System Embeddings [11.873521953539361]
  Deep learning recommendation models (DLRMs) have been widely applied in Internet companies. We propose a GPU-based software cache approaches to dynamically manage the embedding table in the CPU and GPU memory space. Our proposed software cache is efficient in training entire DLRMs on GPU in a synchronized update manner.
  arXiv  Detail & Related papers  (2022-08-08T12:08:05Z)
• NumS: Scalable Array Programming for the Cloud [82.827921577004]
  We present NumS, an array programming library which optimize NumPy-like expressions on task-based distributed systems. This is achieved through a novel scheduler called Load Simulated Hierarchical Scheduling (LSHS) We show that LSHS enhances performance on Ray by decreasing network load by a factor of 2x, requiring 4x less memory, and reducing execution time by 10x on the logistic regression problem.
  arXiv  Detail & Related papers  (2022-06-28T20:13:40Z)
• PLSSVM: A (multi-)GPGPU-accelerated Least Squares Support Vector Machine [68.8204255655161]
  Support Vector Machines (SVMs) are widely used in machine learning. However, even modern and optimized implementations do not scale well for large non-trivial dense data sets on cutting-edge hardware. PLSSVM can be used as a drop-in replacement for an LVM.
  arXiv  Detail & Related papers  (2022-02-25T13:24:23Z)
• TENSILE: A Tensor granularity dynamic GPU memory scheduler method towards multiple dynamic workloads system [9.86589655261934]
  TENSILE is a method of managing GPU memory in tensor granularity to reduce the GPU memory peak. We implement TENSILE on our own deep learning framework, and evaluated its performance.
  arXiv  Detail & Related papers  (2021-05-27T17:46:16Z)
• Nimble: Lightweight and Parallel GPU Task Scheduling for Deep Learning [7.43260596107574]
  We propose Nimble, a deep learning (DL) execution engine that runs tasks in parallel with minimal scheduling overhead. Nable automatically parallelizes the execution of GPU tasks by exploiting multiple GPU streams in a single GPU. evaluation on a variety of neural networks shows that compared to PyTorch, Nimble speeds up inference and training by up to 22.34$times$ and 3.61$times$, respectively.
  arXiv  Detail & Related papers  (2020-12-04T17:25:46Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.