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

TL;DR

DataStates-LLM tackles the bottlenecks of checkpointing at extreme LLM scale by introducing Composable State Providers that expose heterogeneous data structures as a uniform stream. Leveraging the immutability of model/optimizer state during forward/backward passes, it enables lazy, non-blocking captures and overlaps I/O with computation across 3D parallelism. A streaming, multi-level flushing engine backed by kernel-accelerated I/O and a DeepSpeed integration delivers up to $4\times$ checkpoint throughput and up to $2.2\times$ faster end-to-end training, mitigating serialization and heterogeneity bottlenecks. The approach demonstrates strong scalability on 70B-parameter LLMs across 256 GPUs and holds promise for compression-based enhancements and broader shard management in future work.

Abstract

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.

Figures (15)

• Figure 1: Sharding of checkpoints during AI model training for pipeline (PP), tensor (TP), and data (DP) parallelism.
• Figure 2: Aggregate checkpoint sizes of different model sizes and average checkpoint size per GPU.
• Figure 3: Different iteration phases. Model and optimizer states are immutable during forward and backward passes.
• Figure 4: Breakdown of serialization and write performance for different data sizes.
• Figure 5: Overview of DataStates-LLM: Composable state providers capture data structures subject to "3D checkpoint heterogeneity" in a streamlined fashion and flush them to multi-level storage tiers using a checkpoint engine.