Understanding LLM Checkpoint/Restore I/O Strategies and Patterns

TL;DR

This paper investigates I/O strategies for checkpointing large language models, highlighting the importance of aggregation and coalescing to improve throughput and reduce bottlenecks in big-data I/O workflows.

Contribution

It provides a detailed analysis of I/O bottlenecks in LLM checkpointing and proposes microbenchmark-based strategies using liburing to optimize performance.

Findings

Uncoalesced small-buffer I/O halves throughput.

Aggregation and alignment restore bandwidth and reduce metadata overhead.

Our approach outperforms existing engines with up to 7.6x higher throughput.

Abstract

As LLMs and foundation models scale, checkpoint/restore has become a critical pattern for training and inference. With 3D parallelism (tensor, pipeline, data), checkpointing involves many processes, each managing numerous tensors of varying shapes and sizes, that must be persisted frequently to stable storage (e.g., parallel file systems). This turns checkpoint/restore into a big-data I/O problem characterized by volume, variety, and velocity. The workflow must traverse the full storage stack -- from GPU memory through host memory and local storage to external repositories -- whose tiers differ by orders of magnitude in performance, creating bottlenecks under concurrency even with asynchronous flush/prefetch. Kernel-accelerated I/O libraries such as \texttt{liburing} may mitigate these issues versus POSIX, but their effectiveness for LLM checkpointing remains underexplored. We develop microbenchmarks to quantify trade-offs when using \texttt{liburing}, evaluating how aggregation, alignment, and I/O coalescing interact under buffered and direct I/O. We find that uncoalesced small-buffer operations halve throughput relative to synthetic workloads, while file system-aware aggregation restores bandwidth and reduces metadata overhead. Compared to state-of-the-art LLM checkpointing engines, our approach achieves up to $3.9\times$ higher write throughput than DataStates-LLM and $7.6\times$ higher than TorchSnapshot. These results highlight the need for aggregation and coalescing strategies that align with modern file systems and I/O backends.