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

TL;DR

10Cache tackles GPU memory bottlenecks in cloud-scale LLM training by profiling tensor execution order and using a three-tier memory hierarchy (GPU, CPU, NVMe) to prefetch and evict tensors. It introduces a tensor characteristic analyzer, a buffer-reuse cache allocator, a memory-aware tensor allocator, and a prefetch-eviction scheduler that overlap data transfers with computation, enabling efficient placement of FP16 parameters and selective FP32 optimizer states. Experimental results across multiple models show up to a 2× reduction in training time, up to 86.6× GPU cache hit-rate improvements, and up to 2.15× CPU and 1.33× GPU memory utilization gains, demonstrating practical benefits for cloud environments. The approach integrates with DeepSpeed and delivers deployable gains for cost-sensitive, memory-constrained LLM training scenarios.

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.

Figures (20)

• Figure 1: CPU and GPU memory utilization in ZeRO-Infinity.
• Figure 2: LLM model states memory breakdown.
• Figure 3: Tensors timeline for the OPT-125M model. Each dot marks a unique PyTorch operation at which the corresponding tensor becomes active.
• Figure 4: 10Cache's three components' end-to-end flow from profiling a DNN model to tensor allocation across CPU, GPU and NVMe memory. Tensors from the same layer share the same color, while different layers use distinct colors.
• Figure 5: Cache buffer allocation strategy (A contiguous memory block is partitioned into fixed-size buffers for different tensor sizes. Buffers are registered in a free list by size.)