[TEST PR] Add resumable training functionality

[leesharkey]

Description

This PR adds checkpoint-based resume functionality to SPD training jobs, enabling seamless continuation of interrupted
training runs. This is critical for cluster environments where jobs may be killed and rescheduled.

Key Features

• Full state preservation: Saves and restores model weights, optimizer state (momentum buffers), RNG states (PyTorch,
  NumPy, Python, CUDA), and dataloader position
• Hybrid resume modes:
  • Auto-resume: Automatically detects and resumes from the latest checkpoint
  • Explicit resume: Resume from a specific checkpoint path
• Config validation: Errors on breaking changes (architecture differences), warns on non-critical changes
  (hyperparameters)
• WandB integration: Continues the same WandB run when resuming
• Distributed training support: Works with multi-GPU and multi-node setups (DDP)
• Deterministic resume: Achieves bit-exact reproduction of training trajectories

Implementation Details

New module: spd/checkpoint.py

• save_checkpoint(): Saves full training state including all RNG states
• load_checkpoint(): Loads and validates checkpoints with config compatibility checks
• find_latest_checkpoint(): Auto-detects the most recent checkpoint by step number

Modified: spd/configs.py

• Added auto_resume: Enable automatic checkpoint detection
• Added resume_from_checkpoint: Path to explicit checkpoint file
• Added wandb_run_id: For continuing existing WandB runs

Modified: spd/run_spd.py

• Resume detection and checkpoint loading logic
• Conditional faithfulness warmup skip on resume
• Dataloader fast-forward with RNG preservation to maintain exact training sequence
• Fixed dataloader position calculation to account for alive_tracker batch and inclusive step counting

Related Issue

N/A - Feature requested for upcoming cluster deployment

Motivation and Context

We're moving to a cluster setup where training jobs may be killed and rescheduled. Without resume functionality, this would
mean:

• Lost training progress
• Wasted compute resources
• Inability to run long training jobs reliably

This PR enables resilient training that can survive interruptions while maintaining deterministic behavior.

How Has This Been Tested?

Test 1: Constant LR Schedule (TMS)

• Phase 1: Trained steps 0→50, saved checkpoint at step 25
• Phase 2: Resumed from step 25→50
• Result: Loss differences < 0.1% at all checkpoints

Test 2: Cosine LR Schedule (TMS)

• Phase 1: Trained steps 0→100 with cosine schedule + 10% warmup, saved checkpoint at step 50
• Phase 2: Resumed from step 50→100
• Result:
  • All LR values match exactly (bit-perfect)
  • Loss values are identical (0.00% difference)
  • Confirms LR schedule continues correctly after resume

Test Configuration

• 100 training steps, checkpoint at step 50
• Cosine schedule with 10% warmup
• Verified at steps 60, 80, 100

Step   | Phase 1 LR   | Phase 2 LR   | LR Match | Phase 1 Loss   | Phase 2 Loss   | Loss Diff %
------------------------------------------------------------------------------------------------------
60     | 0.000635000  | 0.000635000  | ✓        | 0.00959301740  | 0.00959301740  | 0.00%
80     | 0.000329000  | 0.000329000  | ✓        | 0.01291703433  | 0.01291703433  | 0.00%
100    | -0.000018000 | -0.000018000 | ✓        | 0.01030338556  | 0.01030338556  | 0.00%

Bug Fixes During Development

1. Off-by-one error in dataloader position calculation
2. RNG contamination from dataloader fast-forward
3. RNG contamination from alive_tracker batch consumption
4. Missing numpy import
5. Type checking issues with CUDA RNG state storage

All bugs were identified through testing and fixed before commit.

Does this PR introduce a breaking change?

No breaking changes. This is a purely additive feature:

• All new config fields have default values (auto_resume=False, resume_from_checkpoint=None)
• Existing training runs continue to work exactly as before
• Resume functionality is opt-in via config

Usage Examples

Auto-resume from latest checkpoint

auto_resume: true
save_freq: 1000  # Save every 1000 steps

Resume from specific checkpoint

resume_from_checkpoint: /path/to/model_5000.pth

Continue WandB run

resume_from_checkpoint: /path/to/checkpoint.pth
wandb_run_id: abc123xyz  # ID from previous run