PyTorch Distributed Checkpointing API

Legacy serialization using torch.save operates on a single rank, effectively serializing the entire cluster's memory through one bottleneck. This approach becomes mathematically untenable as model parameters ($P$) and optimizer states ($O$) grow. If $P + O$ exceeds the host memory of Rank 0, the training job crashes. Even if memory is sufficient, the serialization latency scales linearly with model size, leaving GPUs idle for extended periods.

The torch.distributed.checkpoint (DCP) API fundamentally shifts this operation from a centralized gather to a distributed parallel write. DCP allows every rank in the FSDP group to stream its local shard of the state dictionary directly to persistent storage. This reduces the effective write time from $T \propto M$ (total model size) to $T \propto \frac{M}{N}$ (where $N$ is the number of ranks), assuming the storage backend supports sufficient IOPS.

Architecture of Distributed Checkpointing

DCP operates differently than standard PyTorch serialization. Instead of pickling a single Python object, it orchestrates a planned graph of tensor writes. When a save is triggered, DCP creates a metadata file describing the global tensor structure and a series of binary shards containing the actual heavy data.

The storage representation is decoupled from the runtime topology. This separation enables a critical feature: topology-agnostic loading. A model trained on 128 GPUs can be checkpointed and subsequently resumed on 64 GPUs, provided the total memory capacity is sufficient. The DCP loader reads the metadata and re-shards the tensors to match the current process_group configuration.

Data flow in a distributed checkpointing operation where local shards interact with the DCP Planner to generate parallel write streams.

Implementing Sharded Saving

To utilize DCP with FSDP, we must explicitly configure the model to yield a sharded state dictionary. By default, calling .state_dict() on an FSDP module might attempt to gather the full weights, which defeats the purpose of using DCP. We use the FSDP.state_dict_type context manager to enforce SHARDED_STATE_DICT.

The following implementation demonstrates how to save both model weights and optimizer states. Note that the optimizer state must also be handled through the FSDP API to ensure it corresponds correctly to the sharded parameters.

import torch.distributed.checkpoint as dcp
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import StateDictType

def save_checkpoint(model, optimizer, step, checkpoint_path):
    # Ensure path exists on rank 0, or let the writer handle it
    # The StateDictType must be SHARDED for distributed saving
    with FSDP.state_dict_type(model, StateDictType.SHARDED_STATE_DICT):
        # 1. Create the state dictionary
        # This does NOT move data to CPU or gather it to one rank
        state_dict = {
            "model": model.state_dict(),
            # Optimizer must be FSDP-aware
            "optimizer": FSDP.optim_state_dict(model, optimizer),
            "step": step
        }

        # 2. Execute the distributed save
        # distinct from torch.save, dcp.save handles the layout
        dcp.save(
            state_dict=state_dict,
            checkpoint_id=checkpoint_path,
        )

# Usage within training loop
# save_checkpoint(fsdp_model, optimizer, current_step, "checkpoints/step_1000")

In this implementation, dcp.save utilizes a FileSystemWriter by default. Each rank writes its portion of the data to a specific file structure inside checkpoint_path. The overhead is minimal because no cross-node communication is required to aggregate tensors.

Asynchronous Checkpointing

For models exceeding 100 billion parameters, even parallel IO can take seconds or minutes depending on the storage bandwidth. Pausing computation for this duration reduces the Model Flops Utilization (MFU). DCP supports asynchronous saving, allowing the training loop to proceed immediately while IO operations occur on background threads.

To enable this, we typically rely on async_save=True (available in newer PyTorch versions or via snapshotting extensions). However, doing so introduces a race condition: if the training loop updates the weights while the save thread is still reading them, the checkpoint becomes corrupted.

The solution involves capturing a snapshot of the weights in host memory or ensuring the IO completes before the next backward pass modifies the gradients. FSDP creates a copy-on-write mechanism or a fast clone of the sharded references to mitigate this, but explicit synchronization is often safer in custom loops.

def async_checkpoint_handler(model, optimizer, path):
    with FSDP.state_dict_type(model, StateDictType.SHARDED_STATE_DICT):
        state_dict = {
            "model": model.state_dict(),
            "optimizer": FSDP.optim_state_dict(model, optimizer)
        }
        
        # Start the save operation
        # Note: Ensure your storage backend supports concurrent writes efficiently
        future = dcp.async_save(
            state_dict=state_dict,
            checkpoint_id=path
        )
        return future

# In training loop
# future = async_checkpoint_handler(model, opt, "ckpt/step_N")
# ... perform forward pass ...
# future.result() # Ensure save is done before critical sections if needed

Loading and Resharding

Restoring a checkpoint is not as simple as mapping file names to ranks. Because the cluster size might have changed, DCP uses the metadata file to determine which parts of the global tensor belong to the current rank's shard.

When loading into an FSDP module, the module must already be initialized and sharded. The dcp.load function reads the data in place. This is memory efficient because we never materialize the full model; we only read the specific bytes required for the local GPU.

def load_checkpoint(model, optimizer, checkpoint_path):
    # We must use the same StateDictType context
    with FSDP.state_dict_type(model, StateDictType.SHARDED_STATE_DICT):
        # Create a placeholder state_dict with the correct structure
        # The values in this dict serve as the "plan" for loading
        state_dict = {
            "model": model.state_dict(),
            # Optimizer loading requires the model structure to be known
            "optimizer": FSDP.optim_state_dict(model, optimizer),
            "step": 0 # Placeholder
        }

        # Load directly into the placeholders
        dcp.load(
            state_dict=state_dict,
            checkpoint_id=checkpoint_path,
        )
        
        # Apply the loaded optimizer state back to the optimizer engine
        # This is a critical step often missed
        FSDP.optim_load_state_dict(model, optimizer, state_dict["optimizer"])
        
        return state_dict["step"]

The FSDP.optim_load_state_dict function is mandatory. Standard PyTorch optimizers do not understand sharded states natively. FSDP acts as the translator, scattering the loaded optimizer partitions to the correct parameter groups on each device.

Performance Implications

The transition to sharded checkpointing dramatically alters the IO profile of the training cluster. In a legacy setup, network bandwidth on Rank 0 is the limiting factor. With DCP, the limitation shifts to the aggregate write bandwidth of the storage system.

When configuring the distributed file system (e.g., Lustre, GPFS) or Object Store (S3, Azure Blob), it is important to ensure the backend can handle $N$ simultaneous connections.

Comparison of save latency between legacy gathering and Distributed Checkpointing as model size increases.

As illustrated, legacy serialization creates an exponential bottleneck, while DCP maintains a near-linear profile relative to the per-GPU memory shard size rather than the total model size. For terabyte-scale models, DCP is not merely an optimization; it is the only functional mechanism for persistence.