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

arXiv: https://arxiv.org/abs/2601.16956 · HTML: https://mlsanity.com/paper/2601.16956/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 18 of its 47 references that are in the map.
- First posted on arXiv 23 Jan 2026.

## Builds on (most cited first, 18 of 18)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [BLOOM: A 176B-Parameter Open-Access Multilingual Language Model](https://mlsanity.com/paper/2211.05100/) (2022)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://mlsanity.com/paper/1910.02054/) (2019)
- [Mixed Precision Training](https://mlsanity.com/paper/1710.03740/) (2017)
- [An overview of gradient descent optimization algorithms](https://mlsanity.com/paper/1609.04747/) (2016)
- [PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel](https://mlsanity.com/paper/2304.11277/) (2023)
- [ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning](https://mlsanity.com/paper/2104.07857/) (2021)
- [Spike No More: Stabilizing the Pre-training of Large Language Models](https://mlsanity.com/paper/2312.16903/) (2023)
- [Varuna: Scalable, Low-cost Training of Massive Deep Learning Models](https://mlsanity.com/paper/2111.04007/) (2021)
- [Oobleck: Resilient Distributed Training of Large Models Using Pipeline Templates](https://mlsanity.com/paper/2309.08125/) (2023)
- [Unicron: Economizing Self-Healing LLM Training at Scale](https://mlsanity.com/paper/2401.00134/) (2023)
- [DataStates-LLM: Lazy Asynchronous Checkpointing for Large Language Models](https://mlsanity.com/paper/2406.10707/) (2024)
- [Robust LLM Training Infrastructure at ByteDance](https://mlsanity.com/paper/2509.16293/) (2025)
- [TRANSOM: An Efficient Fault-Tolerant System for Training LLMs](https://mlsanity.com/paper/2310.10046/) (2023)
- [FastPersist: Accelerating Model Checkpointing in Deep Learning](https://mlsanity.com/paper/2406.13768/) (2024)
- [Understanding LLM Checkpoint/Restore I/O Strategies and Patterns](https://mlsanity.com/paper/2512.24511/) (2025)

## Built on by (most cited first, 1 of 1)

- [ReCoVer: Resilient LLM Pre-Training System via Fault-Tolerant Collective and Versatile Workload](https://mlsanity.com/paper/2605.11215/) (2026)

## Abstract (arXiv)

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× higher checkpointing throughput and reduces end-to-end training time by up to 2.2× compared to state-of-the-art solutions, effectively mitigating the serialization and heterogeneity bottlenecks in extreme-scale LLM training.