Research map / Large language models / Reasoning and chain of thought
Cost-Efficient LLM Training with Lifetime-Aware Tensor Offloading via GPUDirect Storage
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Teraio is a lifetime-aware tensor offloading framework that profiles tensor lifetimes in PyTorch, plans SSD/host-memory offload and prefetch with an I/O-aware algorithm, and migrates tensors via GPUDirect Storage to expand GPU memory for LLM training.
Problem
LLM training is limited by GPU memory, and existing SSD/host offloading (e.g., ZeRO-Offload, ZeRO-Infinity) at coarse granularity wastes SSD bandwidth and GPU memory and stalls training.
Approach
Offload large inactive tensors at fine granularity, using profiled tensor size and inactive-period lifetimes to schedule offload/prefetch under available SSD bandwidth, preferring SSDs and falling back to host memory when SSD bandwidth saturates, with direct GPU-SSD transfers via GPUDirect Storage.
Key results
- Teraio improves training throughput over ZeRO-Offload and ZeRO-Infinity by 1.47x on average and reaches 80.7% of ideal unlimited-GPU-memory performance.
- On Llama3-70B, Teraio-Mixed reaches 59.7% of ideal performance and outperforms ZeRO-Infinity by 1.33x.
- Teraio-SSD and Teraio-Mixed reduce training cost by 5.88x and 5.41x versus a pure-GPU setup for Llama3-70B.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use lifetime-aware tensor migration | LLM training with TERAIO · Various LLMs and configurations | training throughput vs. ZeRO-Offload and ZeRO-Infinity: 1.47× on average | Helped |
| use TERAIO-Mixed instead of ZeRO-Infinity | Llama3-70B training · Llama3-70B | training throughput: 1.33× | Helped |
| use SSDs and CPU memory for tensor migration (ablation) | TERAIO · Granite-code-base-8B | training throughput vs. TERAIO-SSD | Helped |
| restrict tensor migration to SSDs instead of using SSDs and CPU memory (ablation) | TERAIO · Granite-code-base-8B | training throughput | Hurt |
| use SSDs and CPU memory for tensor migration (ablation) | TERAIO · Llama3-70B | training throughput vs. ideal: 59.7% | Helped |
| increase (micro-)batch size using memory expansion | TERAIO · Llama3-8B | training throughput: 9% | Helped |
| use 2 SSDs and 64 GB of CPU memory (ablation) | TERAIO · Llama-8B | training throughput vs. ideal | Helped |
| use 2 SSDs and 512GB of CPU memory instead of 4 SSDs and 1,024GB of CPU memory (ablation) | TERAIO · Llama-70B | training throughput relative to the larger-resource setup: 73.1% | Hurt |
| migrate tensors only to SSDs rather than using SSDs and CPU memory (ablation) | TERAIO · Llama3-70B training setup | cost savings vs. PureGPUs: 5.88× | Helped |
| use SSDs and CPU memory rather than SSDs only (ablation) | TERAIO · Llama3-70B training setup | cost savings vs. PureGPUs: 5.41× | Helped |
Surprising findings
- Active tensors account for only 1.7% on average of allocated GPU memory per iteration, suggesting most GPU memory holds inactive tensors.
Limitations
- Llama3-70B with Teraio-Mixed still reaches only 59.7% of ideal performance, so higher migration bandwidth is needed for very large models.
- Experiments use only 2 H100 GPUs and 8 SSDs on a single server, limiting scale evaluation.
- Full-precision training only; mixed-precision comparison deemed unfair and not evaluated.
Experimental setup
Datasets: C4
Metrics: training throughput percentage of ideal performance cost efficiency migration bandwidth utilization
Baselines: ZeRO-Offload ZeRO-Infinity Ideal case (unlimited GPU memory) PureGPUs
Builds on: ZeRO-Offload ZeRO-Infinity PyTorch TorchTitan
Closest papers by meaning
- SpeedLoader: An I/O efficient scheme for heterogeneous and distributed LLM operation (NeurIPS 2024 · similarity 0.91)
- Practical Offloading for Fine-Tuning LLM on Commodity GPU via Learned Sparse Projectors (AAAI 2025 · similarity 0.90)
- Speculate Deep and Accurate: Lossless and Training-Free Acceleration for Offloaded LLMs via Substitute Speculative Decoding (NeurIPS 2025 · similarity 0.90)
- 70% Size, 100% Accuracy: Lossless LLM Compression for Efficient GPU Inference via Dynamic-Length Float (DFloat11) (NeurIPS 2025 · similarity 0.88)
- Demystifying Cost-Efficiency in LLM Serving over Heterogeneous GPUs (ICML 2025 · similarity 0.88)
- ReCoVer: Resilient LLM Pre-Training System via Fault-Tolerant Collective and Versatile Workload (NeurIPS 2026 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.