Can I resume training from just the model weights?

Usually not. Reliable resume requires restoring the optimizer, learning rate scheduler, GradScaler, random number generator states, training step, gradient accumulation position, and data loading progress. Otherwise, the learning rate, sample order, or optimization trajectory will change upon resume.

Can I load an old checkpoint directly after changing the number of GPUs?

It depends on the checkpoint format, parallelism strategy, and framework capabilities. PyTorch DCP supports re-sharding within its coverage based on target shard information. Changes involving combinations of TP, PP, ZeRO, etc., require framework-specific universal checkpoints or conversion flows, and compatibility drills should be completed first.

Is asynchronous saving always better than synchronous saving?

Not necessarily. Asynchronous saving reduces the training main-path pause but consumes CPU buffers, pinned memory, and storage bandwidth, and requires managing outstanding tasks. Production environments should typically limit to one concurrent async checkpoint and set guardrails on memory and commit latency.

If I only save model weights, why does the loss still decrease after resume?

A decreasing loss does not prove correct resume. The model may still train after changes to optimizer momentum, learning rate phase, random state, and sample order, but the trajectory has shifted. Production acceptance should check state continuity and short-run comparisons, not just whether the task continues to run.

Distributed LLM Fine-Tuning Checkpoint Resume in Production: Sharded Checkpoints, Atomic Commits, and World Size Changes

Background: A Checkpoint Exists ≠ Training Can Resume

In single-node fine-tuning, torch.save(model.state_dict()) is often sufficient for model export. But with FSDP, ZeRO, Tensor Parallel, or Pipeline Parallel, training state is scattered across multiple ranks, and a checkpoint transforms from “one file” into a set of interdependent shards.

Production incidents rarely involve a complete absence of checkpoints. Instead, you encounter:

Some ranks finish writing, others only write halfway before a node failure;
latest points to a new directory, but the new checkpoint hasn’t been fully committed;
Model weights restore successfully, but optimizer momentum, learning rate, or GradScaler are not restored;
The dataset is re-read from the start of an epoch, causing sample duplication;
The original training used 64 GPUs, but only 48 are available for resume, and old shards cannot be loaded into the new topology;
Async saving backlogs, consuming all host memory, eventually dragging down the training process.

Therefore, the goal of production-grade checkpoint resume isn’t “finding a batch of files.” It’s ensuring a checkpoint simultaneously meets four conditions: state completeness, commit determinability, topology compatibility, and resume verifiability.

PyTorch Distributed Checkpoint (DCP) supports parallel save and load across multiple ranks and can re-shard across different world sizes in supported scenarios based on the loading side’s shard information. It produces multiple files, not a single traditional torch.save file, and loading requires the target model to first allocate state storage and provide shard information.

Core Principle: Treat a Resume Point as a Distributed Transaction

1. Define the Training State Boundary

A resumable checkpoint should cover at least:

Model State: Base model parameters, LoRA or other adapter parameters, trainable heads.
Optimizer State: Adam first and second moments, parameter group settings, ZeRO/FSDP shards.
Scheduler State: Current learning rate, warmup progress, scheduler internal counter.
Numerical State: GradScaler, mixed precision configuration, gradient accumulation micro-step.
Random State: Python, NumPy, PyTorch CPU, each CUDA device, and data sampler RNG.
Progress State: Global step, epoch, consumed samples, data cursor.
Configuration Fingerprint: Model architecture, tokenizer, chat template, dataset version, training parameters, and code version.
Parallel Topology: World size, DP/TP/PP dimensions, ZeRO stage, FSDP sharding strategy.

Hugging Face Accelerate’s checkpoint interface saves the model, optimizer, RNG, and GradScaler, and also allows registering custom stateful objects like learning rate schedulers. If resuming mid-epoch, you also need to restore or skip the already consumed DataLoader batches. This detail is often missed: identical weights do not guarantee the model sees the same next sample.

2. Save in Sharded Format, Avoid Forced Aggregation on the Training Path

In large model training, aggregating all shards into a single full-precision FP32 state dict can cause significant communication, CPU memory, and save-pause overhead. A more sensible primary checkpoint is a Sharded Checkpoint: each rank writes its own shard in parallel, while saving enough metadata for the loading side to reorganize.

Full weight export and fault recovery are two different goals:

Goal	Training Resume Checkpoint	Release Artifact
Primary Concern	Fast write, complete state, re-shardable	Universal format, stable file, easy for inference loading
Write Method	Each rank writes shards in parallel	Aggregate into full weights
Lifecycle	Frequently generated during training, cleaned after resume	Long-term retention, consumed downstream

Don’t force full weight aggregation for every training checkpoint just to make publishing easier. DeepSpeed provides tools to extract FP32 weights from ZeRO checkpoints, but the official documentation warns that this step can be CPU memory-bound, and loading aggregated weights back into the model usually requires re-initializing the DeepSpeed Engine.

3. Use Commit Markers to Distinguish “File Exists” from “Resume Point Valid”

It’s recommended to split a save operation into two phases:

Phase A: Write to a Temporary Prefix

checkpoints/
  step-00012000.inprogress/
    rank-00000.distcp
    rank-00001.distcp
    ...
    metadata.part.json

After each rank finishes writing, record the file size, checksum, and state summary. The coordinating rank only enters the commit phase after confirming all expected shards exist and pass verification.

Phase B: Publish Commit Evidence

Object storage typically shouldn’t rely on directory renames to simulate atomic operations. A more robust engineering approach is to use immutable prefixes, then write a very small COMMITTED file or Manifest, and finally update the latest.json pointer. The recovery side only acknowledges versions with valid commit evidence.

{
  "checkpoint_id": "step-00012000",
  "status": "COMMITTED",
  "global_step": 12000,
  "saved_world_size": 64,
  "parallelism": { "dp": 8, "tp": 8, "pp": 1 },
  "training_state": {
    "epoch": 2,
    "micro_step": 3,
    "consumed_samples": 1572864
  },
  "fingerprints": {
    "model_config": "sha256:...",
    "tokenizer": "sha256:...",
    "dataset_manifest": "sha256:...",
    "training_config": "sha256:..."
  },
  "shards": [
    { "rank": 0, "path": "rank-00000.distcp", "size": 4831838208, "sha256": "..." }
  ]
}

latest.json is just an index, not a source of truth. Even if the pointer is corrupted, the system should be able to scan recent committed Manifests for fallback, rather than selecting the version with the largest directory name.

Engineering Implementation

Save Flow: Freeze Consistent State First, Then Async Persist

PyTorch DCP’s async_save first copies the training state to an internal CPU buffer, then writes to storage asynchronously, reducing blocking time on the GPU training main path. But this isn’t a zero-cost operation: the official documentation notes that CPU memory pressure is roughly proportional to the per-rank checkpoint size and the number of ranks; using Pinned Memory can accelerate staging copies but causes peak pinned memory to remain resident long-term.

A simplified save controller:

from concurrent.futures import Future
import torch.distributed as dist
import torch.distributed.checkpoint as dcp

checkpoint_future: Future | None = None

def save_checkpoint(state: dict, checkpoint_id: str) -> Future:
    global checkpoint_future
    # Avoid piling up multiple async save tasks
    if checkpoint_future is not None:
        checkpoint_future.result()
    checkpoint_future = dcp.async_save(
        state_dict=state,
        checkpoint_id=f"{checkpoint_id}.inprogress",
    )
    return checkpoint_future

def commit_checkpoint(checkpoint_id: str) -> None:
    assert checkpoint_future is not None
    checkpoint_future.result()
    dist.barrier()
    if dist.get_rank() == 0:
        verify_all_shards(checkpoint_id)
        write_manifest(checkpoint_id)
        publish_latest_pointer(checkpoint_id)
    dist.barrier()

Here, verify_all_shards, write_manifest, and publish_latest_pointer are business-side commit protocols, not a complete transaction system automatically provided by DCP. The key constraint is: an incomplete Future cannot be treated as a successful checkpoint, and a directory without a published Manifest cannot be used for resume.

Resume Flow: Compatibility Check Must Precede Loading

After a resume task starts, it’s recommended to follow this order:

Find the most recent committed Manifest.
Verify the Manifest, shard count, sizes, and checksums.
Compare model architecture, tokenizer, dataset, and training configuration fingerprints.
Create a new Device Mesh or parallel topology based on the target GPU count.
Construct and pre-allocate the target model and optimizer state dicts.
Load and re-shard the model and optimizer states.
Restore the Scheduler, GradScaler, RNG, global step, and data cursor.
Execute post-resume validation before allowing the task to enter formal training.

state = build_stateful_training_state(
    model=model,
    optimizer=optimizer,
    scheduler=scheduler,
    scaler=scaler,
)
dcp.load(
    state_dict=state,
    checkpoint_id=selected_checkpoint,
)
restore_rng_state(manifest["rng_state"])
restore_data_cursor(manifest["training_state"])
validate_resume_state(model, optimizer, manifest)

DCP loading is in-place, relying on the target model’s already allocated state dict and shard information. Therefore, the order “read the checkpoint first, then decide the new topology” is wrong; you should first decide the resume resources and shard layout, then perform the load.

World Size Changes: Handle in Three Categories

Category 1: Only Data Parallelism Changes. When the model architecture and supported sharding strategy remain unchanged, DCP can re-shard based on the target side’s shard information. This type of change is the most suitable first-stage goal for elastic recovery.

Category 2: Changes in TP, PP, DP, or ZeRO Combinations. This typically requires a universal checkpoint format or framework-specific conversion. DeepSpeed documentation marks Universal Checkpoints as a still-in-development capability for supporting GPU count changes under 3D Parallelism. The core idea of the Universal Checkpointing paper is to decouple the distributed representation at save time from the unified parameter representation and mapping metadata at load time.

In production, you cannot assume any combination is recoverable just because a feature is named “Universal.” A clear compatibility matrix must be established:

resume_compatibility:
  - from: {dp: 8, tp: 8, pp: 1, zero: 1}
    to: {dp: 6, tp: 8, pp: 1, zero: 1}
    status: verified
  - from: {dp: 8, tp: 8, pp: 1, zero: 1}
    to: {dp: 8, tp: 4, pp: 2, zero: 1}
    status: blocked_until_drill_passes

Category 3: Changes in Model Architecture or State Semantics. For example, adding a trainable layer, changing the optimizer, adjusting the vocabulary, or switching ZeRO stages. These cannot be simply treated as world size changes. They should enter an explicit migration flow; if no migrator exists, you must Fail Closed rather than skipping missing keys and continuing training.

How to Prove “It Really Resumed”

Successful loading only proves deserialization didn’t error. Production acceptance also requires short-run comparisons:

Global step, learning rate, and GradScaler match the save point.
Optimizer parameter groups and momentum tensor counts match.
The data cursor hasn’t returned to the start of an epoch.
The first batch ID after resume is as expected.
The next step’s loss, gradient norm, and parameter update magnitude show no anomalous jumps.
After running several consecutive steps, the metric difference from an uninterrupted control task is within a predefined tolerance.

Distributed GPU training may not achieve bit-exact consistency across all kernels, topologies, and parallel configurations. Therefore, resume acceptance shouldn’t just compare final weight hashes. A more practical goal is to verify state continuity, no missing or duplicate data, and no unexplained trajectory mutations.

Fault Injection Drills: Don’t Wait for the First Incident to Validate Recovery

It’s recommended to periodically execute the following drills in a pre-production environment:

Kill a process while some ranks are writing, verify .inprogress is not selected.
Delete one shard, verify the recovery side rejects loading and falls back to the previous commit point.
Tamper with file content, verify the checksum guardrail activates.
Point latest.json to a non-existent version, verify Manifest scan fallback.
Resume with a different world size, verify the re-sharding compatibility matrix.
Save and resume mid-gradient-accumulation, verify the micro-step doesn’t drift.
Simulate object storage timeout, verify the async save Future doesn’t backlog indefinitely.
After resume, run a short-run comparison on fixed samples, verify loss and data cursor.

A checkpoint system without fault injection is just an unvalidated backup script.

Monitoring Metrics and Alerts

At a minimum, record the following metrics:

Metric	Meaning
`checkpoint_save_duration_seconds`	Duration of a single save operation
`checkpoint_async_staging_seconds`	Duration of the async staging phase
`checkpoint_bytes_total`	Total checkpoint size
`checkpoint_commit_lag_seconds`	Latency from save completion to commit completion
`checkpoint_inprogress_count`	Number of currently incomplete checkpoints
`checkpoint_last_committed_step`	Most recently committed global step
`checkpoint_restore_duration_seconds`	Duration of a restore operation
`checkpoint_restore_failures_total{reason}`	Number of restore failures, categorized by reason
`checkpoint_checksum_failures_total`	Number of checksum failures
`checkpoint_reprocessed_samples_total`	Number of samples reprocessed after resume
`resume_loss_delta`	Loss difference before and after resume
`resume_world_size_change_total`	Number of world size change resumes

Alerts should not just look at “whether a directory was recently created.” They should look at how far the current training progress is from the most recently committed and drill-validated resume point. This is the true Recovery Point Objective (RPO) for the training task.

Applicable Scenarios

This approach is particularly suitable for:

Cross-node, long-running SFT, DPO, and continual pre-training tasks.
Training clusters using preemptible or elastic GPU resources.
Models using FSDP, ZeRO, or hybrid parallelism that cannot be fully saved as a single file.
Tasks where the GPU count may be adjusted during training, but the optimizer trajectory must be preserved.
Enterprise environments with strict requirements for no missing/duplicate samples, training auditability, and cost control.

For small LoRA tasks that complete in a few hours, you don’t need to copy the entire complex system. However, it’s still recommended to save the optimizer, scheduler, RNG, data cursor, and configuration fingerprints, rather than just the adapter weights.

Common Misconceptions

Misconception 1: If the model loads, training can resume. Restoring only model parameters loses optimizer momentum, learning rate progress, and numerical scaling state. The result is closer to “restarting training from old weights” than a true resume.

Misconception 2: The directory with the largest name is the latest checkpoint. A half-written directory can also have the largest step. Validity must be judged by the Manifest and commit marker.

Misconception 3: All ranks writing to the same file is simpler. This easily leads to overwrites, lock contention, and corruption. Distributed training should use the framework’s supported sharded save protocol.

Misconception 4: Async saving can be infinitely concurrent. Each task requires CPU staging and I/O. PyTorch’s official recommendation is that most users should keep only one async request at a time to avoid extra memory pressure.

Misconception 5: Variable world size means all parallelism dimensions are variable. Changes in DP count, TP/PP, and model architecture are different levels of migration and must be validated separately.

Misconception 6: It’s okay to re-shuffle data on resume. Duplicate or skipped samples change the training distribution, especially affecting short datasets, preference data, and data recipes with staged mixing.