10Cache: Heterogeneous Resource-Aware Tensor Caching and Migration for LLM Training

2025-11-17

AI Summary

To address GPU memory capacity limitations in cloud-based large language model (LLM) training—where offloading tensors to CPU or NVMe incurs high migration latency and leads to poor heterogeneous memory utilization—this paper proposes a heterogeneous resource-aware tensor caching and migration system. Our method jointly models tensor execution order and size distribution to design a dynamic prefetching strategy, and introduces an adaptive pinned-memory buffer allocation and reuse mechanism to significantly reduce tensor migration overhead and memory allocation latency. Experimental evaluation across diverse LLM workloads demonstrates up to 2.0× training speedup, an 86.6× improvement in GPU tensor cache hit rate, and 2.15× and 1.33× increases in CPU and GPU memory utilization, respectively. The system effectively balances training efficiency with cloud infrastructure cost.

Technology Category

• Machine Learning: Hardware-aware ML
• Natural Language Processing: (Large) Language Models
• Computer Vision: Large Vision Models

Application Category

• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

Training large language models (LLMs) in the cloud faces growing memory bottlenecks due to the limited capacity and high cost of GPUs. While GPU memory offloading to CPU and NVMe has made large-scale training more feasible, existing approaches suffer from high tensor migration latency and suboptimal device memory utilization, ultimately increasing training time and cloud costs. To address these challenges, we present 10Cache, a resource-aware tensor caching and migration system that accelerates LLM training by intelligently coordinating memory usage across GPU, CPU, and NVMe tiers. 10Cache profiles tensor execution order to construct prefetch policies, allocates memory buffers in pinned memory based on tensor size distributions, and reuses memory buffers to minimize allocation overhead. Designed for cloud-scale deployments, 10Cache improves memory efficiency and reduces reliance on high-end GPUs. Across diverse LLM workloads, it achieves up to 2x speedup in training time, improves GPU cache hit rate by up to 86.6x, and increases CPU/GPU memory utilization by up to 2.15x and 1.33x, respectively, compared to state-of-the-art offloading methods. These results demonstrate that 10Cache is a practical and scalable solution for optimizing LLM training throughput and resource efficiency in cloud environments.

Problem

• Reduces GPU memory bottlenecks in LLM training
• Minimizes tensor migration latency across GPU-CPU-NVMe tiers
• Improves memory utilization and training speed via prefetching

Innovation

• Profiles tensor execution for prefetch policies
• Allocates pinned memory based on tensor sizes
• Reuses memory buffers to minimize allocation overhead