Distributed Checkpointing, in depth: global-offset metadata, the save and load planners, resharding by chunk intersection, deduplication and format conversion

What is on disk

Using FileSystemWriter, a DCP checkpoint is a directory with two kinds of file. Data files are named __{rank}_{n}.distcp: rank-prefixed, numbered, holding the raw bytes of whatever that rank wrote. A single .metadata file describes everything. It holds a Metadata object whose state_dict_metadata maps each fully-qualified key (model.layers.0.attn.wq.weight, optim.state.model.layers.0.attn.wq.weight.exp_avg) to either a BytesStorageMetadata for pickled non-tensor objects or a TensorStorageMetadata with three fields: properties (dtype, layout and similar), size (the global shape) and chunks, a list of ChunkStorageMetadata entries each carrying offsets and sizes. Storage-specific data maps each chunk to the file and byte range that holds it.

You can read the index without loading any tensor data, which makes it the first tool for debugging a checkpoint:

from torch.distributed.checkpoint import FileSystemReader
from torch.distributed.checkpoint.metadata import TensorStorageMetadata

md = FileSystemReader("/ckpt/step_12000").read_metadata()
for key, entry in sorted(md.state_dict_metadata.items()):
    if isinstance(entry, TensorStorageMetadata):
        chunks = [(tuple(ch.offsets), tuple(ch.sizes)) for ch in entry.chunks]
        print(key, tuple(entry.size), entry.properties.dtype, len(chunks), chunks[:2])
    else:
        print(key, "bytes")

Two checks belong in every pipeline that consumes checkpoints. First, for each tensor the chunks should tile the global shape exactly: no gaps, no overlaps. Second, the set of keys should match what the loading code expects; a renamed module shows up here as a key that exists on one side only.

The planner protocol

Both save and load are a short conversation between every rank and a coordinator (rank 0 by default), driven by a planner object. For a save:

1. Set up. Each rank calls set_up_planner with its local state dict. Distributed tensors (DTensor, or the older ShardedTensor) are expanded into their local shards, each tagged with its global offsets.
2. Local plan. create_local_plan lists the write items this rank could contribute: one per tensor shard, one per non-tensor object.
3. Global plan. The coordinator gathers every local plan and runs create_global_plan. This is where replicated items are deduplicated, where files are assigned to ranks, and where the Metadata index is assembled.
4. Finish. Each rank receives its slice of the global plan, may adjust it in finish_plan, and the storage writer writes its items.
5. Commit. Write results flow back to the coordinator, which writes .metadata. A directory without it is an incomplete checkpoint and must be treated as one.

Loading mirrors this. Each rank's local plan lists read items: for every tensor in the state dict it was handed, the intersection of the region it holds with each saved chunk listed in the metadata. That is why loading is always in place: you construct the model in the target layout first, take its state dict, and DCP fills those tensors. The planners are pluggable. DefaultSavePlanner always deduplicates and exposes dedup_save_to_lowest_rank (its older dedup_replicated_tensors argument is deprecated and no longer has any effect); DefaultLoadPlanner exposes allow_partial_load. Custom planners are how frameworks rename keys, transform tensors on the way in, or skip state.

Worked example: resharding by intersection

Make it concrete with a tensor of global shape [12, 4], an embedding table of twelve rows. Two ranks saved it sharded by rows, so the metadata holds chunk A at offsets [0, 0] sizes [6, 4] and chunk B at offsets [6, 0] sizes [6, 4]. Three ranks now load it, also sharded by rows: rank 0 holds rows 0 to 3, rank 1 rows 4 to 7, rank 2 rows 8 to 11.

The planner's work for each pair of (requested region, saved chunk) is a box intersection, done one dimension at a time:

def intersect(req_off, req_size, ch_off, ch_size):
    """Return (offset in chunk, offset in request, lengths), or None if disjoint."""
    in_chunk, in_req, lengths = [], [], []
    for ro, rs, co, cs in zip(req_off, req_size, ch_off, ch_size):
        lo, hi = max(ro, co), min(ro + rs, co + cs)
        if lo >= hi:
            return None
        in_chunk.append(lo - co)
        in_req.append(lo - ro)
        lengths.append(hi - lo)
    return in_chunk, in_req, lengths

chunks = {"A": ([0, 0], [6, 4]), "B": ([6, 0], [6, 4])}
for rank, start in enumerate([0, 4, 8]):
    for name, (co, cs) in chunks.items():
        hit = intersect([start, 0], [4, 4], co, cs)
        if hit:
            print(f"rank {rank} <- chunk {name}: {hit}")

# rank 0 <- chunk A: ([0, 0], [0, 0], [4, 4])
# rank 1 <- chunk A: ([4, 0], [0, 0], [2, 4])
# rank 1 <- chunk B: ([0, 0], [2, 0], [2, 4])
# rank 2 <- chunk B: ([2, 0], [0, 0], [4, 4])

Read the second line: rank 1 takes rows 4 and 5 of the global tensor from offset 4 inside chunk A, and writes them at offset 0 of its own shard. The third line takes the first two rows of chunk B, which are global rows 6 and 7, into offset 2. Nothing in this calculation knows how many ranks saved the checkpoint, which is the whole point.

Two dimensions work identically. Suppose an [8, 8] weight was saved under 2-way FSDP sharding on dimension 0 and 2-way tensor parallelism on dimension 1: four chunks of [4, 4] at offsets [0,0], [0,4], [4,0] and [4,4]. Load it with 4-way tensor parallelism on dimension 1 and no FSDP, and each rank wants an [8, 2] column strip. Rank 0's strip at offsets [0, 0] intersects the two left-hand chunks, so it issues two reads of [4, 2] each. Every rank does the same against its own pair of chunks: eight reads in total, each a strided slice of a saved chunk.

Loading into a different layout

The state-dict helpers in torch.distributed.checkpoint.state_dict handle the awkward part: producing keys that do not depend on how the model is wrapped, and mapping optimizer state from parameter ids to parameter names so it can be resharded like the weights. A save and a load into a different parallel layout look like this:

import torch.distributed.checkpoint as dcp
from torch.distributed.checkpoint.state_dict import (
    get_state_dict, set_state_dict, StateDictOptions)

# --- save, in the training job (any layout) ---
model_sd, optim_sd = get_state_dict(model, optimizer)
dcp.save({"model": model_sd, "optim": optim_sd, "step": step},
         checkpoint_id=f"/ckpt/step_{step}")

# --- load, in a job built with a different mesh ---
model = build_model(new_mesh)            # parameters already sharded for the new layout
optimizer = build_optimizer(model)
model_sd, optim_sd = get_state_dict(model, optimizer)
state = {"model": model_sd, "optim": optim_sd, "step": 0}
dcp.load(state, checkpoint_id="/ckpt/step_12000")   # fills tensors in place
set_state_dict(model, optimizer,
               model_state_dict=state["model"], optim_state_dict=state["optim"])

Three details matter. The optimizer must exist, with its state allocated, before loading, because DCP can only fill tensors that already exist; get_state_dict takes care of this for standard optimizers. Non-tensor values such as step are stored as bytes and returned whole. And StateDictOptions controls the shape of what you get: full_state_dict=True gathers unsharded tensors, with cpu_offload=True to keep them off the GPU, which is what an export path wants, while strict decides whether missing keys are fatal.

For fine-tuning, where the new model has heads the checkpoint lacks, pass a DefaultLoadPlanner(allow_partial_load=True) so keys absent from the checkpoint are left at their initial values instead of failing the load. Keep that option out of resume paths: there, a missing key is a bug you want to hear about.

To leave the DCP world, torch.distributed.checkpoint.format_utils provides dcp_to_torch_save (a DCP directory to one torch.save file) and torch_save_to_dcp for the reverse. Conversion runs in a single process and materialises full tensors, so size the machine for the whole model in host memory.

Replicated tensors and deduplication

When several ranks hold the same tensor (a replicated LayerNorm weight under tensor parallelism, or every parameter under plain data parallelism), each rank's local plan offers it. The default global plan keeps one copy, so a data-parallel checkpoint is not N times the model size. dedup_save_to_lowest_rank changes which rank does the writing; by default each duplicated item goes to whichever candidate rank has the least data assigned so far, which balances bytes written across ranks.

The same rule creates the most damaging DCP bug. A plain tensor saved under one key from several ranks is assumed to be identical on all of them. Per-rank state that genuinely differs, such as RNG states, data-loader positions or per-rank counters, will be deduplicated to a single rank's value, and every rank resumes with it. Give per-rank state per-rank keys (rng/rank_17), or store it outside DCP in small per-rank files.

Failure modes

Storage throughput is the other operational limit. Data transfer paths, including reading from NVMe straight into GPU memory, are covered in the GPUDirect Storage article.

Trade-offs

Pick DCP when the same checkpoint must outlive the layout that wrote it, which for anything trained with FSDP or tensor parallelism is almost always. Keep a converted single-file export for consumers that do not run PyTorch distributed.

What to do next

1. Print the .metadata index of your latest checkpoint and check that each tensor's chunks tile its global shape.
2. Write a CI test that saves a small model on 2 ranks and loads it on 3, then compares every tensor with the original.
3. Audit your state dict for per-rank values stored under shared keys.
4. Confirm the resume path refuses a directory without .metadata.
5. Decide where allow_partial_load is allowed (fine-tuning) and where it is forbidden (resume).
6. Time a dcp_to_torch_save export once, so you know what the export host needs in memory and how long it takes.