DataStates-LLM: Scalable Checkpointing for Transformer Models Using Composable State Providers

Computer Science > Distributed, Parallel, and Cluster Computing [Submitted on 23 Jan 2026] DataStates-LLM: Scalable Checkpointing for Transformer Models Using Composable State Providers Avinash Maurya, M. Mustafa Rafique, Franck Cappello, Bogdan Nicolae The rapid growth of Large Transformer-based models, specifically Large Language Models (LLMs), now scaling to trillions of parameters, has necessitated training across thousands of GPUs using complex hybrid parallelism strategies (e.g., data, tensor, and pipeline parallelism). Checkpointing this massive, distributed state is critical for a wide range of use cases, such as resilience, suspend-resume, investigating undesirable training trajectories, and explaining model evolution. However, existing checkpointing solutions typically treat model state as opaque binary blobs, ignoring the ``3D heterogeneity'' of the underlying data structures--varying by memory location (GPU vs. Host), number of ``logical'' objects sharded and split across multiple files, data types (tensors vs. Python objects), and their serialization requirements. This results in significant runtime overheads due to blocking device-to-host transfers, data-oblivious serialization, and storage I/O contention. In this paper, we introduce DataStates-LLM, a novel checkpointing architecture that leverages State Providers to decouple state abstraction from data movement. DataStates-LLM exploits the immutability of model parameters during the forward and backward passes to perform ``lazy'', non-blocking asynchronous snapshots. By introducing State Providers, we efficiently coalesce fragmented, heterogeneous shards and overlap the serialization of metadata with bulk tensor I/O. We evaluate DataStates-LLM on models up to 70B parameters on 256 A100-40GB GPUs. Our results demonstrate that DataStates-LLM achieves up to 4\times higher checkpointing throughput and reduces end-to-end training time by up to 2.2\times compared to state-of-the-art solutions, effectively mitigating the serialization and heterogeneity bottlenecks in extreme-scale LLM training. Subjects: Distributed, Parallel, and Cluster Computing (cs.DC); Artificial Intelligence (cs.AI); Performance (cs.PF) Cite as: arXiv:2601.16956 [cs.DC]   (or arXiv:2601.16956v1 [cs.DC] for this version)   https://doi.org/10.48550/arXiv.2601.16956 Focus to learn more Submission history From: Avinash Maurya [view email] [v1] Fri, 23 Jan 2026 18:26:14 UTC (713 KB) Access Paper: HTML (experimental) view license Current browse context: cs.DC < prev   |   next > new | recent | 2026-01 Change to browse by: cs cs.AI cs.PF References & Citations NASA ADS Google Scholar Semantic Scholar Export BibTeX Citation Bookmark Bibliographic Tools Bibliographic and Citation Tools Bibliographic Explorer Toggle Bibliographic Explorer (What is the Explorer?) Connected Papers Toggle Connected Papers (What is Connected Papers?) Litmaps Toggle Litmaps (What is Litmaps?) scite.ai Toggle scite Smart Citations (What are Smart Citations?) Code, Data, Media Demos Related Papers About arXivLabs Which authors of this paper are endorsers? | Disable MathJax (What is MathJax?)