MemAscend: System Memory Optimization for SSD-Offloaded LLM Fine-Tuning

Owing to the huge success of generative artificial intelligence (AI), large language models (LLMs) have emerged as a core subclass, underpinning applications such as question answering, text generation, and code completion. While fine-tuning these models on domain-specific data can yield significant performance gains, it also poses daunting computational challenges, especially for researchers and small organizations with limited hardware resources. Although SSD offloading (i.e., ZeRO-Infinity) has emerged as a viable strategy to overcome the GPU memory barrier via leveraging both system memory (i.e., CPU DRAM) and storage space (i.e., solid-state devices, SSDs), its design primarily targets model-centric performance issues. As a result, key system-level issues, including system memory fragmentation, inefficient pinned buffer allocation, peak CPU usage spikes, and file system overhead, remain unaddressed, stifling scalability and inflating costs. Such an observation motivates this paper to introduce MemAscend, a framework that systematically tackles the underexplored system memory bottlenecks in SSD-offloaded LLM training, with a focus on resource-constrained environments. By streamlining pinned-memory allocation, eradicating fragmentation, and mitigating peak overhead, MemAscend reclaims a substantial system memory budget, enabling larger models, longer context windows, and higher batch sizes without exceeding modest hardware limits. Across diverse LLM benchmarks, MemAscend reduces peak system-memory consumption by an average of 55.7% compared with standard SSD offloading techniques, lowering the hardware barrier for fine-tuning and unlocking new possibilities for cost-effective large-scale training on limited-resource machines.

💡 Research Summary

The paper addresses a critical bottleneck that emerges when large language models (LLMs) are fine‑tuned using SSD‑offloading techniques such as ZeRO‑Infinity. While ZeRO‑Infinity successfully moves static model states (weights, optimizer moments) from GPU memory to CPU DRAM and NVMe SSDs, it leaves system‑memory management largely untouched. The authors identify three major sources of inefficiency: (1) severe fragmentation of pinned (cudaHostAlloc) memory, which wastes up to 72.71 % of the buffer pool due to aggressive alignment; (2) redundant CPU‑overflow‑check logic that duplicates temporary buffers and inflates peak memory usage by up to 1.25×; and (3) the reliance on the file system as an intermediary between CPU DRAM and the SSD, adding latency and extra copying overhead. These issues scale with model size, limiting the maximum trainable model, context length, and batch size on machines with modest DRAM (e.g., 128 GiB).

MemAscend is introduced as a comprehensive system‑memory optimization framework that tackles these problems through four tightly integrated components:

1. Adaptive Buffer Pool – Instead of fixed‑size staging buffers, MemAscend dynamically fits offloaded tensors into variable‑size slices, reusing freed space and eliminating fragmentation.

2. Alignment‑Free Pinned Memory Allocation – Custom C++ extensions allocate pinned memory at page granularity (4 KB) without the default 256‑byte alignment, cutting allocation overhead roughly in half and reclaiming on average 55 % of the pinned‑memory pool.

3. Fused Overflow‑Check Mechanism – The traditional two‑stage overflow detection (pre‑ and post‑step) is merged into a single pass, removing duplicated temporary buffers and preventing double‑peak memory spikes.

4. Direct NVMe Engine – By bypassing the OS file system and issuing NVMe commands directly from user space, MemAscend reduces I/O latency, system‑call overhead, and eliminates staging copies.

Additionally, the framework incorporates a half‑precision optimizer that stores parameters, gradients, and momentum in fp16/bf16, further reducing the data transferred to the SSD by 58 % and improving overall throughput by up to 24.21 % (or 56.80 % depending on hardware configuration).

The authors evaluate MemAscend on a range of LLMs from 7 B to 70 B parameters using a single node equipped with 128 GiB of DRAM and a high‑performance NVMe SSD. Compared with vanilla ZeRO‑Infinity, MemAscend achieves an average 55.7 % reduction in peak system‑memory consumption. This memory saving translates into dramatic scalability gains: the same hardware can now handle context windows up to 131 k tokens (versus 16 k tokens originally) and batch sizes up to 32 (versus 4), without exceeding memory limits. I/O volume drops by 58 %, and training throughput improves by up to 24 % on a single‑node setup and by as much as 56.8 % in multi‑GPU configurations. The results demonstrate that system‑memory inefficiencies, not just GPU memory, are the dominant limiting factor in SSD‑offloaded fine‑tuning for resource‑constrained environments.

The paper positions MemAscend as a practical, hardware‑agnostic solution that can be layered on top of existing ZeRO‑Infinity deployments without requiring specialized storage accelerators (unlike Smart‑Infinity) or drastic algorithmic changes (unlike LoHan). By focusing on memory allocation, fragmentation, and I/O pathways, MemAscend lowers the cost barrier for academic labs, startups, and individual researchers who wish to fine‑tune large models on commodity servers. The authors suggest future work on extending the adaptive buffer pool to multi‑node clusters, standardizing direct‑NVMe APIs for broader adoption, and integrating activation‑offloading techniques (e.g., SSDTrain) to achieve end‑to‑end memory minimization across both static and residual states.