Pretraining Engineering: Checkpoints, Fault Tolerance, and Training Stability

Summary: Pretraining is not "getting the code running"—it is running for months on a 10K-GPU cluster without incidents. How do you write checkpoints that preserve progress without slowing training? How do you recover fastest when a node dies? How do you triage an anomalous loss curve? This post gives engineering answers: checkpoint strategy, fault-recovery flow, and a loss-anomaly triage checklist.

0. Experiment Goals

• Design a checkpoint strategy: the period/async/sharding trio
• Master the node-fault recovery flow: from detection to training restart, executable steps
• Build a loss-anomaly triage checklist: data first, then code, then cluster
• Understand training-stability guardrails (watchdog, health checks, log audit)

1. Checkpoint Trio: Period / Async / Sharding

# Checkpoint config in a training framework (Megatron/DeepSpeed-style semantics)
checkpoint = {
    "interval": 500,          # save every 500 steps
    "async_save": True,       # async save: does not block training
    "shard": True,            # sharded storage: each card stores its own shard
    "keep_last": 3,           # keep the last 3 (disk-cost control)
    "save_optimizer": True,   # save optimizer state (else momentum is lost after resume)
}

Step-by-step:

• Period (interval): step interval. Too frequent slows training (sync saves block); too sparse loses progress. Rule of thumb: 500-2000 steps, balanced against hardware-failure probability.
• Async save (async_save): write to local disk first, then copy to shared storage in the background—saving does not block training steps; the standard for large-scale training.
• Sharding: each rank stores only its own weight shard (Megatron's distributed_optimizer), avoiding "one checkpoint file of hundreds of GB, centralized write crashes".
• Saving optimizer state: weights only = optimizer momentum zeroed on resume—the convergence curve regresses. Training-interrupted recovery must save optimizer state.

2. Fault Tolerance: What to Do When a Node Dies

# Fault detection: training-process watchdog (alerts and collects logs on death)
# Common signals: rank heartbeat timeout, NCCL timeout, OOM, node unresponsive

# Recovery flow (executable steps):
# 1) Locate the faulty node: scontrol show job + cluster monitoring (DCGM/IB status)
# 2) Isolate it: slurm DRAIN / K8s cordon, prevent reassignment
# 3) Hardware handling: replace bad GPU/NIC; check IB faults (see KB [IB Admin](/en/guide/03-network/ib-admin-guide))
# 4) Restart training: resume from the latest checkpoint
srun --nodes=N --ntasks-per-node=8 \
  python train.py --resume-from /share/ckpt/latest/

Engineering notes (fault tolerance):

• The checkpoint is the only recovery dependency: restart = resume from the latest checkpoint, so "a corrupted checkpoint" is worse than "a dead node"—verify checkpoint integrity (weight hash) before continuing after recovery.
• Grade faults: single-node fault (most common) → restart and recover; large-scale fault (network partition) → don't rush to recover—find the root cause first, or it will recur after restart.
• Save and restore must be auditable: each checkpoint records step/loss/config hash; the resume log confirms "resumed from step N"—preventing silent resume from an older version.

3. Loss-Anomaly Triage Checklist (In Order)

Symptom	First suspect	Action
Loss spike/divergence (NaN/Inf)	data problem	check recent batches: duplicates, anomalous tokens, cleaning-rule mis-kills
Loss plateaus	LR/ratio	check LR schedule, data mix ratio (Data Pipeline)
Loss rebound after restart	checkpoint missing optimizer state	confirm resume includes optimizer state
Inconsistent loss across GPUs	communication/parallel config	check DDP/TP config, NCCL topology
Occasional NaN	numerical stability	AMP checks, gradient clipping, BF16 overflow

How it works:

• The triage order is iron: data first, then code, then cluster. Pretraining data is TB-scale; one cleaning-rule bug can poison an entire training run (KB Data Pipeline stresses "on loss anomaly, check the data report first").
• Loss rebound after restart is the most common "false anomaly": weights saved but optimizer state not—recovery equals "a different optimizer starting point". Check whether the checkpoint contains optimizer.
• Occasional NaN: in large-model BF16 training, gradient under/overflow is an ongoing risk; loss_scaler (AMP) and gradient clipping (max_grad_norm) are the two lines of defense.

4. Training-Stability Guardrails (Key to Long Runs)

# 1) Watchdog: training-process health probe + auto-restart + log retention
# 2) Metrics audit: record loss/throughput/GPU utilization to a time-series DB every N steps
# 3) Config versioning: hash the training script into logs; verify consistency on resume
# 4) Periodic validation: after each checkpoint save, weight-hash + loadability test
# 5) Circuit breaker: N consecutive anomalous steps → auto-pause + alert, avoiding "sick training" running for days

Engineering takeaways:

1. Checkpoint trio: period 500-2000 steps, async save, sharded storage + optimizer state.
2. Fault recovery flow: locate → isolate → handle → recover; verify the checkpoint after recovery.
3. Triage order is iron: data first, then code, then cluster.
4. Restart rebound ≠ broken training: first check whether the checkpoint includes optimizer state.
5. Circuit breaker: auto-pause on consecutive loss anomalies; don't waste compute on sick training.
6. Config versioning: the resume log confirms "resumed from which step, which config".
7. Health checks in the cluster: watchdog + DCGM/IB probes, auto-DRAIN on faults.

5. Further Reading

• Megatron-LM checkpoint docs
• DeepSpeed checkpointing docs
• KB: 10. Pre-trained Models · Pretraining Data Pipeline · NCCL Tuning · Distributed Parallelism