# 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
