DCP model + optimizer¶

KempnerForge uses torch.distributed.checkpoint (DCP) to save and load model and optimizer state. DCP is designed for sharded state — each rank writes only its local slice, the reader automatically reshards into whatever parallelism the loader has. No “rank 0 gathers everything” step.

Entry point: CheckpointManager.save() in kempnerforge/checkpoint/manager.py.

What goes into the DCP shard¶

optim_state = get_optimizer_state_dict(self.model, self.checkpointed_optimizers)
dcp_state = {
    "model":                   get_model_state_dict(self.model),
    "optimizer":               optim_state,
    "optimizer_never_stepped": _never_stepped_record(optim_state),
}
dcp_state, layout = self._present_experts(dcp_state)
save = self._async_ckpt.save
if layout is not None:
    save = partial(save, planner=DefaultSavePlanner(dedup_save_to_lowest_rank=True))
save(dcp_state, checkpoint_id=str(dcp_dir), process_group=self._process_group)

_present_experts returns the state as it is unless a layer splits its experts with expert parallelism. Then it places each rank’s experts, and their optimizer state, among all of the layer’s experts, so the checkpoint holds every expert; load does the same before DCP and undoes it after. Such a save also writes each chunk that several ranks hold from its lowest rank, so the EP ranks’ copies of the other weights are not mixed. See Expert parallelism § Checkpoints.

checkpointed_optimizers is the optimizer plus any inner optimizers it hands parameters to — see Inner optimizers.

get_model_state_dict / get_optimizer_state_dict are the DCP-aware helpers from torch.distributed.checkpoint.state_dict — not raw model.state_dict() / optimizer.state_dict(). They key the optimizer state by parameter fully-qualified name (not positional index) and keep the FSDP/DTensor sharding intact, which is what makes load (and resharding) line up by name. See Loading for why the raw calls break resume.

Three top-level keys. DCP introspects the state dicts, finds DTensor / ShardedTensor parameters, and writes each shard to disk with enough metadata to reassemble. The third is empty whenever every parameter has optimizer state, and an empty entry writes no key at all, so an ordinary run’s checkpoints keep the two-key shape any reader already expects. A save before the first optimizer step writes the model alone (see below).

What’s in each:

model state — every parameter and buffer: weights, RMSNorm scales, learned RoPE frequencies (if present), and any registered buffer. Under FSDP2 these are DTensors; under TP they’re DTensors on a 2D mesh. DCP handles both.
optimizer state — AdamW’s exp_avg, exp_avg_sq, step counters; Lion’s exp_avg; Muon’s momentum_buffers plus its inner AdamW’s moments. All per-parameter tensors live on the same device and parallelism shape as the parameter, so DCP saves them symmetrically.
never-stepped record — one entry per parameter the optimizer holds without any state, i.e. one that had not yet received a gradient. An optimizer creates per-parameter state lazily, so there is nothing to save for those, and once the checkpoint is written that is indistinguishable from state that went missing. The record keeps the distinction for the load — see Loading.

Before the first optimizer step (the step-0 checkpoint of a dyn_ckpt_window from step 0) none of the checkpointed_optimizers holds state, so the save writes neither the optimizer nor the record. There is nothing to save then, and reading the optimizer would create state: get_optimizer_state_dict first gives an optimizer that holds none a dummy lr=0 step, which leaves AdamW one step ahead in its bias correction and makes ScheduleFreeAdamW divide by zero. A load that finds no optimizer in a checkpoint — this one, or one holding model weights alone — loads it like a pretrained model: the weights and train state, with every checkpointed optimizer reset to its freshly built state. It reads that from the DCP metadata before building any optimizer template, so no optimizer takes the dummy step. Past step 0 it warns, because the run’s optimizer moments are then missing and the resume does not continue it exactly.

Not in the DCP shard: scheduler, dataloader, RNG, and training metadata. Those go in train_state.pt alongside — see Train state.

Inner optimizers¶

get_optimizer_state_dict keeps only an optimizer’s own state and param_groups; any other key its state_dict() returns is dropped. Muon hands every parameter but the hidden nn.Linear weights — the embedding and output head, conv kernels, packed expert weights, MoE router gates unless optimizer.muon_router is set, norms and biases — to an inner torch.optim.AdamW that its own param_groups do not list, so passing Muon alone would leave that AdamW’s moments out of the checkpoint and restart them from zero on every resume.

Muon.inner_optimizers names the inner AdamW, and the manager’s checkpointed_optimizers — [optimizer, *optimizer.inner_optimizers] — go to both the getter and the setter. Their state shares the "optimizer" entry, keyed by parameter name like the rest, with the inner AdamW’s param groups after Muon’s — so a reader that passes Muon alone still finds Muon’s state where it expects it.

A Muon checkpoint saved before the inner AdamW was checkpointed holds only Muon’s own param groups. The load detects that from the DCP metadata, restores Muon alone, leaves the inner AdamW fresh — as every resume from such a checkpoint did before — and logs a warning that the resume is not bit-exact. The next checkpoint it writes is complete.

On-disk layout¶

checkpoints/
├── step_1000/
│   ├── .metadata                         ← DCP global manifest
│   ├── __0_0.distcp                      ← rank-0 model shard
│   ├── __1_0.distcp                      ← rank-1 model shard
│   ├── ...                               ← one `.distcp` per rank
│   ├── train_state.pt                    ← non-distributed state
│   └── metadata.json                     ← human-readable step info
├── step_2000/                            ← same layout
└── latest -> step_2000                   ← symlink

.distcp files are the actual tensor storage (one per rank, written in parallel). .metadata is the global index that lets the reader figure out which .distcp file contains what.

With pipeline parallelism¶

PP makes each stage hold a different set of parameters — DCP needs disjoint shards per stage, so save() writes to a per-stage subdirectory:

# manager.py
dcp_dir = ckpt_dir / f"pp{self._pp_rank}" if self._pp_rank is not None else ckpt_dir

checkpoints/step_1000/
├── pp0/                           ← stage 0 shards (embedding + first layers)
│   ├── .metadata
│   └── __*_0.distcp
├── pp1/                           ← stage 1 shards
│   ├── .metadata
│   └── __*_0.distcp
├── pp2/                           ← ...
├── pp3/                           ← last stage (final norm + output head)
└── train_state.pt                 ← one per checkpoint, written by global rank 0

Each stage also gets a process group of its own ranks: every rank at the same pp coordinate, whatever the other mesh dimensions are (all of the stage’s DP and TP ranks, or the rank alone under pure PP). build_checkpoint_manager in training/entry.py builds it:

ckpt_pg = build_stage_group(device_mesh)
ckpt_mgr = CheckpointManager(config.checkpoint, model, optimizer,
                             process_group=ckpt_pg, pp_rank=pipeline.rank)

build_stage_group (in distributed/pipeline_parallel.py) creates every stage’s group on every rank, in stage order, and returns the one that holds this rank, since torch.distributed.new_group is collective over the default group.

Without a per-stage group, DCP coordinates the save over every rank: it writes a single .metadata, into pp0/, that indexes the shards written into the other stages’ directories, so the save succeeds and the other stages cannot load their checkpoint.

Save modes¶

AsyncCheckpointer wraps DCP’s save / async_save behind a mode flag:

Config	Behavior	Use
`async_mode = "disabled"`	`dcp.save()` — blocking	Simple, debugging, small models
`async_mode = "async"`	`dcp.async_save()` — snapshots to CPU, writes in background	Default for production
`async_mode = "async_with_pinned_mem"`	Async with pinned-memory staging (faster GPU→CPU copy)	Very large models where GPU→CPU throughput bottlenecks the snapshot

Every save() call first does self.wait() — the previous async save must fully complete before starting a new one. This avoids holding two CPU snapshots at once and avoids racing on the same directory.

The returned future is an AsyncCheckpointerFuture; wait() calls .result() which blocks until the background writer has flushed to disk. Training calls ckpt_mgr.wait() once before shutdown (training/loop.py::run_training_loop) to flush any pending save.

Process groups¶

The process_group= kwarg on dcp.save and dcp.load scopes the all-gather / all-reduce calls that DCP uses internally. Rules:

Non-PP: use the default global group (process_group=None). Every rank holds a slice of the same state.
PP: use a per-stage group. Stage i’s ranks (all DP × TP ranks at PP position i) coordinate alone — they produce one DCP shard set under pp{i}/.

CheckpointManager stores the group at construction time and passes it through on every save/load.

Loading¶

Load is the mirror image of save:

dcp_state = {
    "model":     get_model_state_dict(self.model),
    "optimizer": get_optimizer_state_dict(self.model, self.checkpointed_optimizers),
}
dcp_state, layout = self._present_experts(dcp_state, dcp_dir, saved_keys)
dcp.load(dcp_state, checkpoint_id=str(dcp_dir), process_group=self._process_group)
if layout is not None:
    dcp_state = layout.to_local(dcp_state)
set_model_state_dict(self.model, dcp_state["model"])
set_optimizer_state_dict(
    self.model, self.checkpointed_optimizers, optim_state_dict=dcp_state["optimizer"]
)

Under expert parallelism _present_experts places each rank’s experts among all of them before dcp.load, as on save, and layout.to_local gives the setters this rank’s names and tensors back. For any model with MoE layers it first refuses a checkpoint that holds fewer experts than the model. See Expert parallelism § Checkpoints.

A Muon checkpoint saved before the inner AdamW was checkpointed is the one exception: the load passes Muon alone — see Inner optimizers.

The getter call gives DCP the shape to fill — it doesn’t contain the saved data, just the tensor metadata (and, crucially, the allocated optimizer moment tensors) DCP needs to know what to load where. dcp.load mutates those tensors in place with the loaded values; the setters then write them back into the live model and optimizer.

Why the DCP-aware helpers, not optimizer.state_dict()? On resume the optimizer is freshly built, so optimizer.state_dict() has no exp_avg / exp_avg_sq tensors yet — AdamW creates the per-parameter state lazily on the first .step(). Passing that empty dict as the load template gives dcp.load nothing to fill, so the saved moments are silently dropped and Adam momentum resets to zero at every resume (a non-bit-exact resume). get_optimizer_state_dict allocates the moment tensors up front in the right sharded layout, so dcp.load repopulates them. The model side would work with either call — its parameters are always allocated — but we use the matching getter/setter for symmetry.

The template allocates state for every trainable parameter, while the checkpoint has none for the parameters recorded as never stepped. A load planner drops exactly those from the template and gives them empty state afterwards, so a run in which part of the model has not yet been reached resumes. Every other state the template asks for stays required: a parameter with missing or partial state, or one the saved optimizer did not hold, fails the load.

A checkpoint written before the record existed carries no such evidence, and an absent state cannot be told apart from one that went missing, so a never-stepped parameter in one of those still fails the load. Start from it in a new checkpoint.dir, with [checkpoint].load_path pointing at it and exclude_from_loading = ["optimizer"], which restores the weights and begins with a fresh optimizer (["optimizer", "train_state"] also starts the step, schedule and data order fresh). In the old checkpoint.dir the resume would take precedence over load_path.

A parameter whose optimizer changed since the save holds another optimizer’s state. An embedding or output head that Muon orthogonalized before it routed parameters by module has a saved momentum_buffer where its AdamW now keeps exp_avg, exp_avg_sq and step. The same holds for any other parameter that change moved, such as a vision tower’s position table or an MQA key projection, for a router gate when optimizer.muon_router changes, and for every parameter when optimizer.name changes. The load compares the template’s state entries with the checkpoint’s metadata before reading anything and fails naming those parameters, instead of leaving dcp.load to stop at the first missing key. Continuing from the weights takes the same warm start, from a checkpoint.dir that holds no checkpoint yet. The comparison covers only the optimizers the load restores, so a Muon checkpoint saved before the inner AdamW was checkpointed is not checked: a parameter that moved to the inner AdamW starts fresh, and one that moved to Muon still fails inside dcp.load.

Each rank runs these checks on its own part of the checkpoint, so under pipeline parallelism one stage can fail them while the others pass. A load that restores optimizer state therefore ends its checks with one all-reduce on the manager’s gloo group, which every rank enters: a rank whose check failed raises its error, and every other rank raises a RuntimeError naming the ranks that failed, before any rank loads anything.

The template also holds every param-group key the optimizer has now, and DCP requires each one. An optimizer whose param groups gained a key after its checkpoints were first written lists it in LEGACY_GROUP_DEFAULTS, with the value those checkpoints were trained with: Muon lists orthogonalize as "local", nesterov as False and adjust_lr as "original", the update it made before those options existed. Where a checkpoint lacks a listed key, the load drops it from the template, the optimizer fills in the listed value as it loads the groups, and the manager logs it. A checkpoint that lacks any other key still fails the load.

Loading with a different GPU count triggers DCP’s automatic resharding — see Resharding.

To skip loading the optimizer (e.g. for fine-tuning), pass exclude_keys=["optimizer"]:

ckpt_mgr.load(path=..., exclude_keys=["optimizer"])   # scripts/eval.py does this