Corrupted, Truncated and Incompatible Checkpoint Diagnosis
Agent-Assisted Engineering Check
Prerequisites
Day 55 Checkpoints and error handling mechanics.
Learning Objective
Diagnose partial-write truncations (EOFError, pickle UnpicklingError), verify SHA256 checksums, and automate fallback to the most recent healthy historical checkpoint.
Instructions
1. Understand failure modes when processes are killed during standard write operations: file size is non-zero, but end of file is truncated.
2. Attempting to torch.load a truncated checkpoint throws _pickle.UnpicklingError: pickle data was truncated or EOFError.
3. Implement a checksum verification protocol: along with 'ckpt.pt', save 'ckpt.pt.sha256'. Verify hashes before loading.
4. Build an automated multi-stage fallback: if checkpoint_step_10000 fails validation, seamlessly fall back to checkpoint_step_9000 and notify the training supervisor.
Practice & Prediction Log
Write a recovery loader safe_load_checkpoint(checkpoint_dir) that iterates backwards through sorted checkpoint snapshots, rejects corrupted or truncated files, and successfully loads the latest valid state.
Iterate snapshots in reverse sorted order (sorted(files, reverse=True)). If a .sha256 file exists, verify the hash first; then wrap torch.load in try/except (EOFError, RuntimeError, Exception) and skip any file that fails.
import os, torch, hashlib
def safe_load_checkpoint(ckpt_dir):
files = sorted([f for f in os.listdir(ckpt_dir) if f.startswith('ckpt_') and f.endswith('.pt')], reverse=True)
for fname in files:
path = os.path.join(ckpt_dir, fname)
try:
if os.path.exists(path + '.sha256'):
with open(path, 'rb') as f:
digest = hashlib.sha256(f.read()).hexdigest()
with open(path + '.sha256') as f:
if f.read().split()[0] != digest:
raise ValueError('sha256 mismatch')
ckpt = torch.load(path, map_location='cpu')
print(f'Successfully recovered from {fname}')
return ckpt
except Exception as e:
print(f'Skipping corrupted checkpoint {fname}: {e}')
raise RuntimeError('No valid checkpoint found in directory')
Demonstrate automatic fallback to an intact historical snapshot when the newest snapshot is intentionally truncated.
Authoritative Source Resources & Provenance
PyTorch Tutorial: Saving and Loading Models
Relevant Section / Range: Saving & Loading a General Checkpoint for Inference and/or Resuming Training
Verified Version & Date: PyTorch 2.x Documentation (Verified: 2026-10-30)
Official documentation detailing state_dict serialization, optimizer state saving, and safe loading patterns.
Python, PyTorch tensors and nn.Module basics.
All checkpoint serialization semantics, dictionary schema keys, and safe recovery steps are self-contained in lesson code.
Compared Candidate Sources & Selection Analysis (2)
TensorFlow trackable checkpoint architecture; PyTorch's explicit dictionary is our target implementation standard.
Safetensors provides safe zero-copy tensor storage without pickle, compared directly when discussing pickle execution risks.
nanoGPT: train.py pretraining loop
Relevant Section / Range: train.py: checkpoint dict {'model', 'optimizer', 'model_args', 'iter_num', 'best_val_loss', 'config'} written with torch.save to out_dir/ckpt.pt; the init_from == 'resume' branch reloads it
Verified Version & Date: Git master commit 71c3905 (Verified: 2026-10-30)
Compact historical code reference for the named implementation functions; CS336 spring2025 provides the lead component exercises. Do not infer top-p, atomic persistence or cloud budget enforcement from nanoGPT.
Transformer architecture from Days 46-48 and training loop from Days 49-51.
Complete checkpoint saving patterns and resuming loops are directly implemented in lesson exercises.
Compared Candidate Sources & Selection Analysis (2)
High-complexity multi-GPU distributed model-parallel checkpointing; nanoGPT provides accessible single-GPU clarity.
High-level framework callback abstraction; nanoGPT keeps explicit control over state dict dictionaries.