Checkpoint Storage and Fast Recovery: A Two-Tier Design from Local NVMe to Parallel File System

Note: This is a review article (no first-hand experimental data), based on public practice and architecture analysis.

Experiment Objective

At the day-scale of LLM training, a single unexpected interruption (node crash, power jitter, accidental OOM kill) can destroy more than 10 hours of progress. Checkpoints are the only "undo button," but their storage design is usually ignored: save too often and you stall training (IO bursts eat bandwidth); save too rarely and you lose progress (recovery means retraining); store only on local disk and you lose it (node dies, disk dies with it); store only on a parallel file system and it is too slow and expensive. The goal: a practical two-tier checkpoint architecture (local NVMe hot tier + parallel file system cold tier), plus methods to size save frequency, dual-write timing, and the recovery path.

Code Walkthrough

Part 1: Save-frequency control in the training loop

# Training-loop snippet: checkpoint by token, not by step
import time

ckpt_interval_tokens = 50_000_000   # save every 50M tokens
tokens_since_ckpt = 0
last_ckpt_time = time.time()

for step, batch in enumerate(train_loader):
    loss = train_step(batch)
    tokens_since_ckpt += batch.num_tokens()
    # Strategy 1: step-based saving is unstable (step length varies with data);
    # token-based is more uniform.
    if tokens_since_ckpt >= ckpt_interval_tokens:
        save_checkpoint(step, local_path='/mnt/local-nvme/ckpt-latest',
                        async_copy_to='/mnt/pfs/ckpt/step-%d' % step)
        tokens_since_ckpt = 0
        # Record save duration as the IO-budget gauge
        elapsed = time.time() - last_ckpt_time
        print(f"[ckpt] step={step} save+copy={elapsed:.1f}s")
        last_ckpt_time = time.time()

Key points:

1. Token-based is more stable than step-based: with data parallelism, tokens per step fluctuate with padding and batch alignment. Token-based saving keeps the worst-case lost progress (here, 50M tokens) constant.
2. Write to local NVMe first, copy to PFS asynchronously: local NVMe sequential writes reach GB/s, finishing in tens of seconds; the PFS copy is asynchronous and does not block the training thread (the next part shows how).
3. Print the save duration: this is the IO-budget dashboard. If save+copy exceeds a few minutes, checkpointing is slowing training — adjust frequency or storage tier.

Part 2: Two-tier storage and async copy

# Tier 1 on each node (local NVMe, mounted before training)
mkdir -p /mnt/local-nvme/ckpt-latest && chmod 1777 /mnt/local-nvme

# Tier 2: parallel file system directory (visible to all nodes)
mkdir -p /mnt/pfs/ckpt && chmod 1777 /mnt/pfs/ckpt

# Async dual-write: local first, PFS copy in background (simplified)
import threading, shutil

def async_copy(src, dst):
    threading.Thread(target=lambda: (shutil.copytree(src, dst, dirs_exist_ok=True)), daemon=True).start()

save_checkpoint(step, '/mnt/local-nvme/ckpt-latest')
async_copy('/mnt/local-nvme/ckpt-latest', f'/mnt/pfs/ckpt/step-{step}')

Design logic:

• Tier 1 (local NVMe) guarantees "saving never blocks training" — the local write finishes in tens of seconds and training continues. The risk is data loss on node failure, which is why it is only a buffer.
• Tier 2 (PFS) guarantees "recoverable across nodes" — the PFS is cluster-shared, and any node can read the latest checkpoint. The PFS copy is slower (hundreds of MB/s to GB/s) but asynchronous.
• The tiers connect through async tasks: the training process does not wait for the PFS copy; it only guarantees the previous copy completed before the next save (in practice, a queue plus watermark confirmation).

Principles

Why checkpoints are IO-burst workloads

Average training IO is small (reading data streams, writing logs), but saving a checkpoint creates a large, concentrated write burst: for a 70B model in BF16, weights plus optimizer state (Adam stores two momentum copies) plus data-parallel gradient stats can reach 300GB-1TB. Two problems follow:

1. Bursts saturate shared storage bandwidth: if all nodes save simultaneously (step-aligned), the PFS is saturated instantly and read latency for other jobs spikes — this is why "when training saves, the whole cluster stalls."
2. Write time is not negligible: at 300GB and 500MB/s per PFS client, that is 10 minutes. If saving is synchronous and training saves every 30 minutes, it pauses 10 of those — MFU drops 25%.

Why the two-tier design is right

Tier	Medium	Latency/bandwidth	Reliability	Role
1	Local NVMe	multi-GB/s	lost on node failure	hot: saving does not block training
2	Parallel FS	100MB/s-GB/s	replicas/erasure	cold: recoverable across nodes

The core tradeoff: training continuity vs recovery completeness. Tier 1 keeps training moving; tier 2 keeps recovery possible. In production, tier 2 may even copy every Nth tier-1 checkpoint — dropping one or two intermediate checkpoints in exchange for never being slowed by IO.

Optimizing the recovery path

Recovery cost = reading the checkpoint + rebuilding distributed state. Three levers:

1. Staged recovery: restore weights first (evaluation/inference within seconds), restore optimizer state asynchronously (training waits a little longer), decoupling "can recover" from "fully recovered."
2. Prefer the nearest tier-1 copy: when the same node survives, read directly from local NVMe for second-level recovery; fall back to PFS only when the node is gone.
3. Async save + fixed frequency: fixed token intervals with async dual-write keep the worst-case loss constant — you never face the dilemma of "saving too often slows training, too rarely loses too much."

Results

Estimates for a 70B training scenario (estimated, for design reference):

Parameter	Estimate	Note
Checkpoint size	300-400GB	140GB weights + 160GB+ optimizer state
Local NVMe write	~2-4 min	multi-GB/s sequential
PFS async copy	10-15 min	~500MB/s
Recovery (local)	< 5 min	weights first, optimizer state later
Recovery (PFS)	15-30 min	full read + distributed-state rebuild

Engineering Takeaways

1. Set frequency by tokens: decide the acceptable maximum lost progress first, then derive the save interval — do not save every 500 steps on a whim.
2. Local disk is a necessity, not a luxury: tier 1's cost buys an interruption-free training run; far better value than synchronous PFS writes.
3. Turn the IO budget into a monitored item: save duration and PFS queue watermark belong on the Grafana board (a storage panel next to the DCGM panel from the previous article); catch burst anomalies early.
4. Coordinate with the scheduler: the checkpoint-complete timestamp should be visible to the scheduler, so SLURM/K8s only allow scaling down or node maintenance after a checkpoint — avoiding "saved for nothing."

Next article preview: this one solved "how to come back fast after interruption"; the inference-side memory pressure is the next battlefield — "KV Cache Optimization and Memory Management: An Accounting Method for Long-Context Inference."