docs: checkpoints + eval-diag spec + plans (v1 superseded by v2)

2026-05-31 16:57:43 +02:00
parent 7b3309edcc
commit 13d8ed76da
3 changed files with 3173 additions and 0 deletions
--- a/docs/superpowers/plans/2026-05-31-checkpoints-and-eval-diag-v2.md
+++ b/docs/superpowers/plans/2026-05-31-checkpoints-and-eval-diag-v2.md
--- a/docs/superpowers/plans/2026-05-31-checkpoints-and-eval-diag.md
+++ b/docs/superpowers/plans/2026-05-31-checkpoints-and-eval-diag.md
--- a/docs/superpowers/specs/2026-05-31-checkpoints-and-eval-diag-design.md
+++ b/docs/superpowers/specs/2026-05-31-checkpoints-and-eval-diag-design.md
@@ -0,0 +1,310 @@
+# IntegratedTrainer Checkpoints + Eval-Phase Diag Emission — Design
+
+**Date:** 2026-05-31
+**Status:** Draft (pre-implementation)
+**Branch target:** `ml-alpha-checkpoints-eval-diag` (off `ml-alpha-regime-observer`)
+**Related work:** F4.1 (just shipped), regime-observer F1–F5
+**Pearls / feedback referenced:**
+- `feedback_no_atomicadd`
+- `feedback_no_partial_refactor`
+- `feedback_wire_everything_up`
+- `feedback_single_source_of_truth_no_duplicates`
+- `pearl_adaptive_carryover_discipline`
+
+---
+
+## 1. Motivation
+
+Two gaps surfaced during F4.1 cluster validation:
+
+1. **No checkpoints.** Long runs (20k train + 5k eval ≈ 1h) crash and lose everything. PVC ENOSPC (alpha-rl-tvdpv @ step 10k) wasted ~30 min of GPU. Beyond crash recovery, we need checkpoints for **inference deployment** — the trained model has zero way to be served outside the trainer.
+2. **Eval phase emits no per-step diag.** `step_with_lobsim_gpu` (the GPU-only fast path used by eval) bypasses `DiagStaging`, so `diag.jsonl` is silent during the 5k-step eval window. Monitoring eval health is currently impossible until `eval_summary.json` lands at the end.
+
+Both must be fixed before the next cluster submission. Doing them together because they share the same persistence-layer boilerplate (CudaSlice download, bincode serialization, file I/O).
+
+---
+
+## 2. Goals
+
+### G1 — Inference checkpoints (lite)
+Periodically save weights-only checkpoints sufficient to run `step_with_lobsim_gpu` for inference. No optimizer state, no replay, no ISV.
+
+### G2 — Resume checkpoints (full)
+Periodically save full trainer state sufficient to resume training mid-run with identical trajectory (modulo non-determinism we already accept). Includes weights + Adam state + ISV bus + popart stats. Replay buffer is **optional** (separate flag, can be skipped to save disk).
+
+### G3 — Eval-phase per-step diag JSONL
+Eval emits one JSONL line per step to `eval_diag.jsonl`, with the same schema as `diag.jsonl`. Train and eval go to separate files (no step-namespace confusion).
+
+### G4 — CLI surface
+- `--checkpoint-every <N>` (default `5000`) — save full checkpoint every N train steps
+- `--checkpoint-keep <K>` (default `3`) — rolling window of K most recent full checkpoints
+- `--inference-ckpt-every <N>` (default `5000`) — save lite inference checkpoint every N train steps
+- `--resume-from <path>` — load checkpoint at startup; trainer state restored, training continues from saved `step`
+- `--eval-diag-jsonl <path>` (default `<out>/eval_diag.jsonl`) — eval-phase per-step diag output
+
+### G5 — Non-goals
+- Checkpoint format compatibility across SHA. If trainer struct changes, old checkpoints become unloadable. We accept this; cluster runs are pinned to a SHA anyway. Per `feedback_no_legacy_aliases` — no migration shims.
+- Distributed/multi-node checkpoints. Single L40S / H100 only.
+- Streaming checkpoint upload to MinIO. Local file only; PVC retention handles persistence.
+
+---
+
+## 3. Design overview
+
+### 3.1 Two checkpoint variants
+
+| Field | Lite (inference) | Full (resume) |
+|---|---|---|
+| Encoder weights (`CfcTrunk`) | ✓ | ✓ |
+| Head weights (`DqnHead`, `IqnHead`, `DuelingQHead`, `PolicyHead`, `ValueHead`) | ✓ | ✓ |
+| Target-net weights (`w_target_d`, `b_target_d`, IQN/Dueling target variants) | ✓ | ✓ |
+| Adam optimizer moments (per-head `m_d`, `v_d`, `t`) | ✗ | ✓ |
+| ISV bus snapshot (1024 floats — controller EMAs, popart stats, regime state) | ✗ | ✓ |
+| Trainer scalars (`last_pi_loss`, `last_q_loss`, `last_v_loss`, train step count) | ✗ | ✓ |
+| Replay buffer (PER tree + samples) | ✗ | optional via `--checkpoint-replay` |
+
+**Why split:** A lite checkpoint is ~10-50 MB (weights only). A full checkpoint is ~hundreds of MB (Adam moments roughly double weights, replay buffer at b_size=1024, capacity=32768 is GB-scale). Inference deployments need to be small. Resume needs to be complete.
+
+### 3.2 File layout per run
+
+Under `--out <dir>` (= existing JSONL out):
+```
+diag.jsonl                                      ← train per-step diag (existing)
+eval_diag.jsonl                                 ← NEW: eval per-step diag
+alpha_rl_train_summary.json                     ← existing
+eval_summary.json                               ← existing
+checkpoints/
+    train_step_5000.full.ckpt                   ← rolling, last K kept
+    train_step_10000.full.ckpt
+    train_step_15000.full.ckpt
+    inference_step_5000.lite.ckpt               ← rolling, last K kept
+    inference_step_10000.lite.ckpt
+    inference_step_15000.lite.ckpt
+    inference_final.lite.ckpt                   ← written once after train phase
+    inference_post_eval.lite.ckpt               ← written once after eval phase
+```
+
+Rolling-window deletion happens after each successful save (delete oldest beyond K).
+
+### 3.3 Serialization
+
+`bincode::serialize` to flat binary, matching the existing `CfcTrunk::save_checkpoint` pattern (`feedback_single_source_of_truth_no_duplicates` — don't introduce a second format).
+
+Top-level struct:
+```rust
+struct IntegratedCheckpoint {
+    /// Format version. Bumped on any breaking layout change.
+    /// Loader validates `version == EXPECTED_VERSION` or returns Err.
+    version: u32,
+
+    /// Git SHA of the trainer that produced this checkpoint.
+    /// Logged on load for archaeology; mismatch does NOT block load
+    /// (responsibility is on the caller to know what they're loading).
+    sha: String,
+
+    /// Train step at which this checkpoint was taken (0 for inference
+    /// checkpoints saved before training).
+    step: u64,
+
+    /// Wall-clock UTC timestamp of save, RFC3339. Diagnostic only.
+    saved_at: String,
+
+    /// "lite" or "full". Loader switches on this string.
+    kind: String,
+
+    /// Encoder weights — delegate to existing CfcTrunk::Checkpoint.
+    encoder: CfcTrunkCheckpoint,
+
+    /// Per-head weights (always present).
+    heads: HeadWeightsCheckpoint,
+
+    /// Adam moments + step counters per head. Present only when kind="full".
+    adam: Option<AdamStateCheckpoint>,
+
+    /// ISV bus snapshot — 1024 floats. Present only when kind="full".
+    isv: Option<Vec<f32>>,
+
+    /// Trainer scalar bookkeeping. Present only when kind="full".
+    scalars: Option<TrainerScalarsCheckpoint>,
+
+    /// Replay buffer snapshot. Present only when kind="full" AND
+    /// caller passed `--checkpoint-replay`.
+    replay: Option<ReplayBufferCheckpoint>,
+}
+```
+
+`HeadWeightsCheckpoint` enumerates every CudaSlice<f32> in each head:
+- `DqnHead`: w_d, b_d, w_target_d, b_target_d
+- `IqnHead`: w_embed_d, b_embed_d, w_out_d, b_out_d, w_embed_target_d, b_embed_target_d, w_out_target_d, b_out_target_d
+- `DuelingQHead`: w_v_d, b_v_d, w_a_d, b_a_d, target variants
+- `PolicyHead`: w_d, b_d
+- `ValueHead`: w_d, b_d (verify exact fields when writing)
+
+`AdamStateCheckpoint` enumerates every (m, v, t) triple. AdamW save/load methods already exist (per task #32) — use them.
+
+`ReplayBufferCheckpoint` matches the in-memory layout of the PER buffer (priority tree + per-sample features/actions/rewards/next-features/dones, b_size × capacity).
+
+### 3.4 Public API on `IntegratedTrainer`
+
+```rust
+impl IntegratedTrainer {
+    /// Save a lite (weights-only) checkpoint.
+    pub fn save_inference_checkpoint(&self, path: &Path) -> Result<()>;
+
+    /// Save a full (resumable) checkpoint. `include_replay=true` adds
+    /// the PER buffer (~GBs); `false` skips it (replay rebuilds during
+    /// resume training).
+    pub fn save_full_checkpoint(&self, path: &Path, include_replay: bool) -> Result<()>;
+
+    /// Load a checkpoint (any kind). Loader switches on `kind` field:
+    /// - "lite": loads weights only; trainer's optimizer/ISV stay at
+    ///   the values established by `with_controllers_bootstrapped`.
+    /// - "full": loads weights + Adam + ISV + scalars (+ replay if
+    ///   present). Returns the `step` field so caller can resume the
+    ///   training loop at the right index.
+    pub fn load_checkpoint(&mut self, path: &Path) -> Result<u64>;
+}
+```
+
+### 3.5 Eval-phase diag emission
+
+The current train loop calls `trainer.step_with_lobsim(...)` which writes to `DiagStaging::write_record()` → flushes to `diag.jsonl`. Eval uses `step_with_lobsim_gpu` which has no diag hook.
+
+**Two design choices considered:**
+
+**(A) Add diag emission inside `step_with_lobsim_gpu`** — symmetric, but doubles writes for any non-eval caller that doesn't want diag (none currently). Requires plumbing a `DiagStaging` ref through the GPU path.
+
+**(B) Expose `trainer.snapshot_diag_record() -> DiagRecord` and have the eval loop call it explicitly.** Caller decides destination file. Cleaner separation; no plumbing through hot GPU path.
+
+**Choose B.** The eval loop in `alpha_rl_train.rs:1384-1399` already runs at the caller side — adding `eval_diag_staging.write_record(trainer.snapshot_diag_record(eval_step)?)` after the `step_with_lobsim_gpu` call is one line. The trainer doesn't need to know about diag staging at all.
+
+The `DiagRecord` struct (currently inline in `DiagStaging::write_record` via serde_json) will be extracted to a typed struct so both train and eval paths produce identical schema. Per `feedback_single_source_of_truth_no_duplicates` — one schema definition, two call sites.
+
+### 3.6 CLI integration in `alpha_rl_train.rs`
+
+New flags (under `Cli` struct):
+```rust
+#[arg(long, default_value_t = 5000)]
+checkpoint_every: usize,                       // 0 disables
+#[arg(long, default_value_t = 3)]
+checkpoint_keep: usize,                        // rolling window
+#[arg(long)]
+checkpoint_replay: bool,                       // include PER in full ckpt
+#[arg(long, default_value_t = 5000)]
+inference_ckpt_every: usize,                   // 0 disables
+#[arg(long)]
+resume_from: Option<PathBuf>,                  // load at startup
+#[arg(long)]
+eval_diag_jsonl: Option<PathBuf>,              // default <out>/eval_diag.jsonl
+```
+
+Train loop modifications:
+1. **Startup**: if `--resume-from` set, call `trainer.load_checkpoint()` BEFORE the train loop. Resume with `start_step = loaded_step + 1`. Skip first `start_step` iterations of the data loader (or seek the loader to that position).
+2. **Inside train loop**: every `checkpoint_every` steps, call `save_full_checkpoint()`. Every `inference_ckpt_every` steps, call `save_inference_checkpoint()`. Both write to `<out>/checkpoints/` with rolling deletion.
+3. **End of train phase**: write `<out>/checkpoints/inference_final.lite.ckpt`.
+4. **End of eval phase**: write `<out>/checkpoints/inference_post_eval.lite.ckpt`.
+
+Eval loop modifications:
+1. Open `eval_diag.jsonl` writer at eval phase start.
+2. After each `step_with_lobsim_gpu` call, write a diag line via `trainer.snapshot_diag_record(eval_step)`.
+
+### 3.7 Resume semantics
+
+`--resume-from` loads the checkpoint and continues at `step + 1`. Resume rebuilds:
+- Encoder, all heads, target nets — from saved weights
+- Adam moments — from saved state (each head's optimizer step continues at saved `t`)
+- ISV bus — bit-exact restoration of all 1024 slots (controller EMAs, popart stats, regime observer state)
+- Trainer scalars — `last_*_loss`, popart counters, cumulative_dones
+
+Open replay handling: if checkpoint has replay, use it. If not, PER rebuilds from the resumed training steps (1024 samples per step → 32 steps to refill at capacity=32768). For mid-run crashes the replay loss is bounded by ~32 steps of "cold start" data which is acceptable.
+
+Random-state determinism: scoped_init_seed already pins Xavier init. Mid-run RNG (PER sampling, dropout if any) is NOT saved; resume runs are NOT bit-exact reproductions, only trajectory-similar. Per `pearl_scoped_init_seed_for_reproducibility` — seed is for init only; mid-run RNG drift is accepted.
+
+---
+
+## 4. Implementation phases
+
+### Phase 1 — Schema + lite save/load (smallest unit, no production dependency)
+1. Define `IntegratedCheckpoint` + nested structs in `crates/ml-alpha/src/trainer/checkpoint.rs` (NEW file).
+2. Implement `HeadWeightsCheckpoint::from_trainer(&IntegratedTrainer)` — pure download path.
+3. Implement `IntegratedTrainer::save_inference_checkpoint` (calls trunk's existing save + emits head weights).
+4. Implement `IntegratedTrainer::load_checkpoint` (lite branch only — weights only).
+5. Unit test: round-trip save → load → forward produces identical logits (bit-exact float compare with tolerance ε=1e-6).
+
+### Phase 2 — Full save/load (Adam + ISV + scalars)
+1. Extend struct with `AdamStateCheckpoint`, `TrainerScalarsCheckpoint`.
+2. Implement `save_full_checkpoint(include_replay=false)`.
+3. Implement load_checkpoint full branch (sets ISV, Adam state).
+4. Unit test: save mid-training → load → 1 step → identical losses to non-checkpointed run (modulo non-determinism, ε=1e-3 on loss values).
+
+### Phase 3 — Replay buffer save/load
+1. Add `ReplayBufferCheckpoint` for the PER structure (priority tree + samples + indices).
+2. Implement `--checkpoint-replay` flag handling.
+3. Unit test: save → load → step samples produces identical PER sample indices (deterministic if RNG is seeded).
+
+### Phase 4 — Eval diag emission
+1. Extract `DiagRecord` struct from inline JSON to typed struct in `crates/ml-alpha/src/trainer/diag_record.rs`.
+2. Implement `trainer.snapshot_diag_record(step: usize) -> DiagRecord`.
+3. Update train loop in `alpha_rl_train.rs` to call this instead of inline JSON building.
+4. Add eval loop call site that writes `DiagRecord` to `eval_diag.jsonl`.
+5. Smoke test: run with `--n-eval-steps 100`, verify both `diag.jsonl` (1000 lines) and `eval_diag.jsonl` (100 lines).
+
+### Phase 5 — Resume integration
+1. Add `--resume-from`, `--checkpoint-every`, `--checkpoint-keep`, `--inference-ckpt-every`, `--checkpoint-replay`, `--eval-diag-jsonl` to `Cli`.
+2. Wire resume path in train loop entry.
+3. Wire checkpoint saves in train loop (full at `checkpoint_every`, lite at `inference_ckpt_every`).
+4. Wire rolling-window cleanup.
+5. Wire end-of-phase checkpoints (`inference_final.lite.ckpt`, `inference_post_eval.lite.ckpt`).
+6. Smoke test: kill mid-train, resume from latest full checkpoint, verify trajectory continues without divergence (qpa/l_q/l_pi within ε of where they'd have been).
+
+### Phase 6 — Cluster validation
+1. Submit alpha-rl with checkpoint flags @ b=1024, fold 1, 20k+5k.
+2. Verify checkpoints land in `<out>/checkpoints/`.
+3. Verify eval_diag.jsonl emits per-step entries during eval.
+4. Verify `inference_post_eval.lite.ckpt` is loadable in a separate process (smoke "deployment" test: load → step_with_lobsim_gpu on 1k held-out → produces sane stats).
+
+---
+
+## 5. Disk-space accounting
+
+At b=1024, per-capacity=32768:
+- **Lite checkpoint**: encoder (~50 MB from `CfcTrunk::Checkpoint`) + heads (~5 MB) ≈ **55 MB**
+- **Full checkpoint without replay**: lite + Adam moments (~10 MB) + ISV (4 KB) + scalars (~100 B) ≈ **65 MB**
+- **Full checkpoint with replay**: full + PER (b_size × capacity × per-sample ≈ 1024 × 32768 × ~256 B) ≈ **8 GB** (warning: keeping 3 of these = 24 GB)
+
+Default config (`--checkpoint-keep 3`, no `--checkpoint-replay`):
+- 3 × full (no replay) + 3 × lite + 2 × end-of-phase lite ≈ 290 MB per run
+- vs current `diag.jsonl` at ~3 GB per run → checkpoints are NEGLIGIBLE compared to diag verbosity
+
+With `--checkpoint-replay`: ~24 GB rolling. Need ~30 GB PVC headroom. Cluster runs should NOT default to this — only enable for crash-recovery scenarios where exact replay is needed.
+
+---
+
+## 6. Open decisions for review
+
+| # | Question | Recommendation |
+|---|----------|---------------|
+| 1 | Checkpoint format: bincode vs MessagePack vs HDF5 | **bincode** — matches existing `CfcTrunk::save_checkpoint`, no new deps |
+| 2 | Schema versioning policy | u32 version field, hard-fail on mismatch; no migration shims |
+| 3 | Eval diag: same file or separate | **separate** (`eval_diag.jsonl`) — cleaner step-namespace |
+| 4 | Replay buffer default: included or not | **NOT included** by default — disk savings, replay rebuilds quickly |
+| 5 | Inference checkpoint cadence | match `--checkpoint-every`, configurable separately for flexibility |
+| 6 | Resume behavior on checkpoint/code mismatch | Log SHA mismatch warning, attempt load anyway, hard-fail on struct deserialization error |
+
+---
+
+## 7. Validation criteria
+
+- Phase 1 unit test passes
+- Phase 2 unit test passes
+- Phase 4 smoke produces both diag files with matching schema
+- Phase 5 resume smoke: trajectory continues without divergence
+- Phase 6 cluster run: checkpoints land, eval_diag emits 5000 lines, post-eval lite ckpt loadable in separate process
+
+## 8. Done means
+
+- `alpha_rl_train.rs` accepts all new flags
+- All 5 phase unit tests pass on local RTX 3050
+- One cluster run completes producing checkpoints + eval_diag.jsonl
+- One separate-process inference smoke loads `inference_post_eval.lite.ckpt` and runs forward pass
+- Pearl saved: `pearl_checkpoint_two_variant_lite_full.md` documenting the design rationale