Day 56: Week 8 Consolidation: Optimizer to Checkpoint, plus Resource Planning

Inspect Reference Solution & Trace:

1. Per block = 12 * 768^2 + 4 * 768 = 7,080,960; 12 blocks = 84,971,520. Embeddings = 50,257 * 768 + 1,024 * 768 = 39,383,808. Final LayerNorm = 2 * 768 = 1536. Total N = 124,356,864. Adding the linear biases (12 * 9 * 768 = 82,944) gives 124,439,808, the count PyTorch reports for GPT-2 with bias=True. FP16 weights for N = 124,356,864: 248,713,728 bytes = 237.2 MiB (248.7 MB).
2. Static = 1e9 * 16 bytes = 16 GB. Total = 16 + 8 (activations) + 1 (CUDA context) = 25 GB > 24 GB, so it runs out of memory unless you use gradient checkpointing, a smaller micro-batch with gradient accumulation, or ZeRO offload.
3. C = 6 * N * D = 6 * 125e6 * 10e9 = 7.5e18 FLOPs. Effective throughput = 1.248e14 FLOP/s. Time = 60,096 s = 16.69 hours. Cost = 16.69 * $1.50 = $25.04.

Inspect Reference Solution & Trace:

1. m1 = 0.1 * 0.4 = 0.04, m_hat = 0.04 / 0.1 = 0.4; v1 = 0.001 * 0.16 = 0.00016, v_hat = 0.16, sqrt = 0.4. w1 = 2.0 * (1 - 0.01 * 0.1) - 0.01 * 0.4 / 0.4 = 1.998 - 0.01 = 1.988.
2. Decay applies to tensors with dim >= 2: W^Q and the token embedding. The LayerNorm gain and the MLP bias are 1-D and get no decay (Day 50).
3. Step 50: 5e-4. Step 100: 1e-3. Step 600: 1e-4 + 0.5 * 9e-4 = 5.5e-4. Step 1,100: 1e-4. Step 2,000: 1e-4 (held at the floor).
4. The factor is 1.0 / 4.0 = 0.25. Order: optimizer.zero_grad(set_to_none=True); logits, loss = model(x, y); loss.backward(); clip_grad_norm_(model.parameters(), 1.0); optimizer.step() (Day 51).
5. Logits / 0.5 = [4, 2, 0], so softmax = [0.8668, 0.1173, 0.0159]. Top-p 0.9: token A alone has 0.8668 < 0.9, and A + B reach 0.9841, so keep {A, B} and renormalize to [0.8808, 0.1192, 0] (Day 52).

Inspect Reference Solution & Trace:

1. PPL = exp(2.0794) ≈ 8.0. Validation rises from 2.0794 to 2.31 while training keeps falling, so the run overfits after step 3,000; keep the step-3,000 checkpoint, chosen by validation loss (Day 53).
2. Shuffling tokens puts neighbouring words of the same sentence into both splits, so validation leaks. Split by document first, then tokenize each split (Day 53).
3. A repetition loop: greedy decoding feeds its own output back and falls into an attractor (Day 54). Bigrams: ab, ba, ab, ba, ab = 5 total, 2 unique, so Distinct-2 = 0.4.
4. Without the optimizer state, AdamW restarts with m = v = 0 and step = 0, and without iter_num the schedule jumps back to warmup, so the first updates after resume are badly scaled (Day 55). Payload with atomic write:

payload = {
    'model': model.state_dict(),
    'optimizer': optimizer.state_dict(),   # AdamW m, v and step per parameter
    'iter_num': iter_num,                  # drives get_lr(it) on resume
    'best_val_loss': best_val_loss,
    'rng': torch.get_rng_state(),
    'config': config,
}
torch.save(payload, path + '.tmp')
with open(path + '.tmp', 'rb') as f:
    os.fsync(f.fileno())
os.replace(path + '.tmp', path)

5. Weights 124,356,864 * 4 = 497,427,456 bytes; moments 124,356,864 * 8 = 994,854,912 bytes; total 1,492,282,368 bytes ≈ 1.49 GB (1.39 GiB).

Inspect Reference Solution & Trace:

1. 2.5e9 / 524,288 = 4,768.4, so 4,769 steps; 2% warmup ≈ 95 steps. Halfway through the cosine, lr = 6e-5 + 0.5 * (6e-4 - 6e-5) = 3.3e-4.
2. 124,356,864 * 16 = 1,989,709,824 bytes ≈ 1.99 GB, far below 40 GB. Activations limit the micro-batch, which is why the plan uses 16 x 1,024 with 32 accumulation steps.
3. C = 6 * 124,356,864 * 2.5e9 = 1.865e18 FLOPs. Time = 1.865e18 / 1.092e14 = 17,082 s ≈ 4.74 hours. Cost ≈ 4.74 * $1.20 = $5.69. These are predictions: replace the 35% MFU with the measured value after the first steps and recompute.
4. Run estimate_loss (model.eval(), torch.no_grad(), about 200 batches per split) every 250 steps and at the end; keep the checkpoint with the lowest validation loss, not the last one; each checkpoint holds model, optimizer, iter_num, best_val_loss, RNG state and config, written with tmp + fsync + os.replace; report the validation loss with its perplexity exp(val_loss), on a split made by document.
5. Example: "Samples show local syntax and phrasing learned from the training corpus. A 124M model trained on 2.5B tokens cannot verify facts or reason reliably, and we make no such claim."