Day 55: Checkpoints, State Dictionaries & Safe Resumption

01 120 min

Model State Dictionaries and PyTorch Tensor Serialization

Agent-Assisted Exploration

Prerequisites

Completion of Day 54 Transformer assessment and PyTorch Module fundamentals.

Learning Objective

Understand the exact data structures inside model.state_dict(), how parameters and buffers are distinguished, and how CPU/CUDA mapping works during torch.save/torch.load.

Instructions

Study source: within this activity's existing reading/practice time, use CS336 Assignment 1: Basics (Spring 2025 tagged release), tests/test_serialization.py: test_checkpointing; model, optimizer and iteration state, then explain one connection to the local task in your own words. Keep the local exercise and acceptance checks; use its original explanation if the external source is unavailable. 1. Inspect model.state_dict(). Notice keys are strings corresponding to submodule hierarchies (e.g., 'blocks.0.attn.c_attn.weight').
2. Note that torch.save uses Python pickle under the hood, storing tensor metadata and storage pointers.
3. Compare saving the entire model object vs saving state_dict(). Recognize why official PyTorch guidelines mandate state_dict for code decoupling and architecture stability.
4. Observe map_location='cpu' to prevent OOM errors when loading a GPU-trained checkpoint onto a CPU or a different GPU index.

Practice & Prediction Log

Instantiate a tiny 2-layer GPT model. Save its state_dict to 'ckpt_model.pt'. Load it back into a freshly initialized model using load_state_dict() with map_location='cpu'. Verify with torch.allclose() that all weights match to machine precision.

Mark as Completed

Independent Mastery Criterion

Successfully serialize and reload model weights with 100% parameter equality across instances.

Authoritative Source Resources & Provenance

linking MIT

CS336 Assignment 1: Basics (Spring 2025 tagged release)

Author: Stanford University CS336 Staff | Original Source Link ↗

Relevant Section / Range: tests/test_serialization.py: test_checkpointing; model, optimizer and iteration state

Verified Version & Date: spring2025 tag; README, tests and MIT LICENSE checked 2026-11-01

Selection Rationale:

Lead assignment-driven source for pretraining component practice and tests. nanoGPT remains a compact historical code-reading companion, not the current recommended training curriculum.

Source Prerequisites:

Fluent Python, PyTorch, and linear algebra

Safe Fallback (Offline Resilience):

Local self-contained pretraining scripts run on standard local CPUs or budget GPUs without Stanford cluster dependency.

Compared Candidate Sources & Selection Analysis (2)

Karpathy nanoGPT ↗ compared_alternative [MIT]

Simpler code layout, useful complementary reference for minimal PyTorch Transformer.

Hugging Face Transformers Course ↗ compared_alternative [Apache-2.0]

High-level API focus rather than building from raw tensor fundamentals.

linking BSD-3-Clause

PyTorch Tutorial: Saving and Loading Models

Author: PyTorch Documentation Team | Original Source Link ↗

Relevant Section / Range: What is a state_dict?

Verified Version & Date: PyTorch 2.x Documentation (Verified: 2026-10-30)

Selection Rationale:

Official documentation detailing state_dict serialization, optimizer state saving, and safe loading patterns.

Source Prerequisites:

Python, PyTorch tensors and nn.Module basics.

Safe Fallback (Offline Resilience):

All checkpoint serialization semantics, dictionary schema keys, and safe recovery steps are self-contained in lesson code.

Compared Candidate Sources & Selection Analysis (2)

TensorFlow tf.train.Checkpoint Guide ↗ compared_alternative [Apache 2.0]

TensorFlow trackable checkpoint architecture; PyTorch's explicit dictionary is our target implementation standard.

Hugging Face Safetensors Documentation ↗ compared_alternative [Apache 2.0]

Safetensors provides safe zero-copy tensor storage without pickle, compared directly when discussing pickle execution risks.

linking MIT

nanoGPT: train.py pretraining loop

Author: Andrej Karpathy | Original Source Link ↗

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)

Selection Rationale:

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.

Source Prerequisites:

Transformer architecture from Days 46-48 and training loop from Days 49-51.

Safe Fallback (Offline Resilience):

Complete checkpoint saving patterns and resuming loops are directly implemented in lesson exercises.

Compared Candidate Sources & Selection Analysis (2)

Megatron-LM Checkpoint Utility ↗ compared_alternative [BSD-3-Clause]

High-complexity multi-GPU distributed model-parallel checkpointing; nanoGPT provides accessible single-GPU clarity.

PyTorch Lightning ModelCheckpoint ↗ compared_alternative [Apache 2.0]

High-level framework callback abstraction; nanoGPT keeps explicit control over state dict dictionaries.

02 120 min

Optimizer State Dictionaries, Learning Rate Schedulers & RNG Continuity

Independent Mastery Check (no AI assistance).

Prerequisites

Day 55 Activity 1 and Adam optimizer formulation.

Learning Objective

Persist Adam momentum (first and second moments) and step counters, along with RNG generator states, ensuring training restarts without loss spikes.

Instructions

1. Inspect optimizer.state_dict(). Observe the 'state' dictionary mapping parameter IDs to 'step', 'exp_avg' (1st moment), and 'exp_avg_sq' (2nd moment).
2. Understand why omitting optimizer state causes massive loss spikes upon resuming (Adam defaults momentum to zero, causing erratic updates).
3. Capture torch.get_rng_state(), torch.cuda.get_rng_state_all(), and python random/numpy states in the checkpoint payload.
4. Build a comprehensive dictionary containing model, optimizer, scheduler, step, best_val_loss, and rng_states.

Practice & Prediction Log

Train a model for 5 steps with AdamW. Save a full checkpoint dict. Resume into a new model and optimizer from that dict, and verify that step 6 produces the exact same loss and gradient update as step 6 in an un-interrupted 10-step reference run.

Mark as Completed

Independent Mastery Criterion

Demonstrate zero loss discrepancy between resumed training and continuous baseline execution.

Authoritative Source Resources & Provenance

linking BSD-3-Clause

PyTorch Tutorial: Saving and Loading Models

Author: PyTorch Documentation Team | Original Source Link ↗

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)

Selection Rationale:

Official documentation detailing state_dict serialization, optimizer state saving, and safe loading patterns.

Source Prerequisites:

Python, PyTorch tensors and nn.Module basics.

Safe Fallback (Offline Resilience):

All checkpoint serialization semantics, dictionary schema keys, and safe recovery steps are self-contained in lesson code.

Compared Candidate Sources & Selection Analysis (2)

TensorFlow tf.train.Checkpoint Guide ↗ compared_alternative [Apache 2.0]

TensorFlow trackable checkpoint architecture; PyTorch's explicit dictionary is our target implementation standard.

Hugging Face Safetensors Documentation ↗ compared_alternative [Apache 2.0]

Safetensors provides safe zero-copy tensor storage without pickle, compared directly when discussing pickle execution risks.

linking MIT

nanoGPT: train.py pretraining loop

Author: Andrej Karpathy | Original Source Link ↗

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)

Selection Rationale:

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.

Source Prerequisites:

Transformer architecture from Days 46-48 and training loop from Days 49-51.

Safe Fallback (Offline Resilience):

Complete checkpoint saving patterns and resuming loops are directly implemented in lesson exercises.

Compared Candidate Sources & Selection Analysis (2)

Megatron-LM Checkpoint Utility ↗ compared_alternative [BSD-3-Clause]

High-complexity multi-GPU distributed model-parallel checkpointing; nanoGPT provides accessible single-GPU clarity.

PyTorch Lightning ModelCheckpoint ↗ compared_alternative [Apache 2.0]

High-level framework callback abstraction; nanoGPT keeps explicit control over state dict dictionaries.

03 120 min

Atomic Checkpoint Persistence, File Locks & Version Schema Protection

Agent-Assisted Engineering Check

Prerequisites

Day 55 Activities 1 and 2, file system atomic operations.

Learning Objective

Implement an atomic write pattern (saving to a temporary file, flushing, and os.replace) and embed format version headers to guard against partial-write corruption upon system crash.

Instructions

1. Understand the danger of directly calling torch.save(..., 'ckpt.pt'): if the host OS is preempted mid-write, 'ckpt.pt' becomes corrupted and unrecoverable.
2. Study atomic renaming via os.replace: write to 'ckpt.pt.tmp', sync to disk, then perform an atomic rename.
3. Embed a explicit 'version': 1 header and 'checksum' or SHA256 signature to guarantee schema integrity.
4. Handle safe rotation: keep last N checkpoints and remove stale checkpoints only after the new checkpoint is confirmed atomic.

Practice & Prediction Log

Write a Python function save_checkpoint_atomic(payload, filepath) that creates a temporary sibling file, writes the payload, flushes to disk, and replaces the target filepath atomically. Test it with a mock signal interruption.

Mark as Completed

Independent Mastery Criterion

Write an atomic checkpoint persistence function that guarantees no partially-written target files on disk.

Authoritative Source Resources & Provenance

linking BSD-3-Clause

PyTorch Tutorial: Saving and Loading Models

Author: PyTorch Documentation Team | Original Source Link ↗

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)

Selection Rationale:

Official documentation detailing state_dict serialization, optimizer state saving, and safe loading patterns.

Source Prerequisites:

Python, PyTorch tensors and nn.Module basics.

Safe Fallback (Offline Resilience):

All checkpoint serialization semantics, dictionary schema keys, and safe recovery steps are self-contained in lesson code.

Compared Candidate Sources & Selection Analysis (2)

TensorFlow tf.train.Checkpoint Guide ↗ compared_alternative [Apache 2.0]

TensorFlow trackable checkpoint architecture; PyTorch's explicit dictionary is our target implementation standard.

Hugging Face Safetensors Documentation ↗ compared_alternative [Apache 2.0]

Safetensors provides safe zero-copy tensor storage without pickle, compared directly when discussing pickle execution risks.

linking MIT

nanoGPT: train.py pretraining loop

Author: Andrej Karpathy | Original Source Link ↗

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)

Selection Rationale:

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.

Source Prerequisites:

Transformer architecture from Days 46-48 and training loop from Days 49-51.

Safe Fallback (Offline Resilience):

Complete checkpoint saving patterns and resuming loops are directly implemented in lesson exercises.

Compared Candidate Sources & Selection Analysis (2)

Megatron-LM Checkpoint Utility ↗ compared_alternative [BSD-3-Clause]

High-complexity multi-GPU distributed model-parallel checkpointing; nanoGPT provides accessible single-GPU clarity.

PyTorch Lightning ModelCheckpoint ↗ compared_alternative [Apache 2.0]

High-level framework callback abstraction; nanoGPT keeps explicit control over state dict dictionaries.

04 120 min

Assessment: Resilient Local Checkpoint Manager and Continuity Verification

Independent Mastery Check (no AI assistance).

Prerequisites

Day 55 Activities 1-3.

Learning Objective

Construct and independently test a production CheckpointManager class featuring atomic save, state verification, and uninterrupted loss curve equivalence.

Instructions

1. Encapsulate checkpoint handling in CheckpointManager(checkpoint_dir, max_to_keep=3).
2. Implement save(step, model, optimizer, rng_state) and load_latest(model, optimizer): load_latest finds the newest valid snapshot and restores it, and save keeps at most max_to_keep files.
3. Set up an end-to-end verification script: train a synthetic LM for 20 steps; at step 10, save a checkpoint and kill the process; spawn a fresh process, load step 10, and train to step 20.
4. Assert that step 20 loss in resumed training matches step 20 loss in an uninterrupted 20-step baseline run within 1e-6 tolerance.

Practice & Prediction Log

Run your CheckpointManager verification test. Confirm that max_to_keep pruning operates correctly and that the loss sequence [step 11..20] is identical to the reference sequence.

Mark as Completed

Independent Mastery Criterion

Demonstrate a complete CheckpointManager class passing unit tests with 100% loss trajectory match.

Authoritative Source Resources & Provenance

linking BSD-3-Clause

PyTorch Tutorial: Saving and Loading Models

Author: PyTorch Documentation Team | Original Source Link ↗

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)

Selection Rationale:

Official documentation detailing state_dict serialization, optimizer state saving, and safe loading patterns.

Source Prerequisites:

Python, PyTorch tensors and nn.Module basics.

Safe Fallback (Offline Resilience):

All checkpoint serialization semantics, dictionary schema keys, and safe recovery steps are self-contained in lesson code.

Compared Candidate Sources & Selection Analysis (2)

TensorFlow tf.train.Checkpoint Guide ↗ compared_alternative [Apache 2.0]

TensorFlow trackable checkpoint architecture; PyTorch's explicit dictionary is our target implementation standard.

Hugging Face Safetensors Documentation ↗ compared_alternative [Apache 2.0]

Safetensors provides safe zero-copy tensor storage without pickle, compared directly when discussing pickle execution risks.

linking MIT

nanoGPT: train.py pretraining loop

Author: Andrej Karpathy | Original Source Link ↗

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)

Selection Rationale:

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.

Source Prerequisites:

Transformer architecture from Days 46-48 and training loop from Days 49-51.

Safe Fallback (Offline Resilience):

Complete checkpoint saving patterns and resuming loops are directly implemented in lesson exercises.

Compared Candidate Sources & Selection Analysis (2)

Megatron-LM Checkpoint Utility ↗ compared_alternative [BSD-3-Clause]

High-complexity multi-GPU distributed model-parallel checkpointing; nanoGPT provides accessible single-GPU clarity.

PyTorch Lightning ModelCheckpoint ↗ compared_alternative [Apache 2.0]

High-level framework callback abstraction; nanoGPT keeps explicit control over state dict dictionaries.