diff --git 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 new file mode 100644 index 000000000..b0e49ba99 --- /dev/null +++ b/docs/superpowers/plans/2026-05-31-checkpoints-and-eval-diag-v2.md @@ -0,0 +1,1251 @@ +# Checkpoints + Eval Diag — Implementation Plan v2 + +> **Supersedes:** `2026-05-31-checkpoints-and-eval-diag.md` (v1). v1 had placeholders, an invented field-name bug, a save-during-training race, and TODO'd the data-loader cursor. v2 has done the pre-flight verification, fixed those, and split into three independent phases. + +> **For agentic workers:** REQUIRED SUB-SKILL: superpowers:subagent-driven-development. Steps use checkbox `- [ ]` syntax. + +**Spec:** `docs/superpowers/specs/2026-05-31-checkpoints-and-eval-diag-design.md` + +**Branch:** `ml-alpha-checkpoints-eval-diag` (off `ml-alpha-regime-observer`) + +**Pre-flight findings (verified against current source):** +- `DuelingQHead` target-net fields: `w_v_target_d, b_v_target_d, w_a_target_d, b_a_target_d` (confirmed at `dueling_q.rs:137-140`) +- `IntegratedTrainer` Adam catalogue (22 instances, all `pub`): + - `dqn_w_adam`, `dqn_b_adam` + - `iqn_w_embed_adam`, `iqn_b_embed_adam`, `iqn_w_out_adam`, `iqn_b_out_adam` + - `policy_w_adam`, `policy_b_adam` + - `value_w_adam`, `value_b_adam` + - `dueling_q_w_v_adam`, `dueling_q_b_v_adam`, `dueling_q_w_a_adam`, `dueling_q_b_a_adam` + - `frd_w1_adam`, `frd_b1_adam`, `frd_w2_adam`, `frd_b2_adam` + - `noisy_mu_w_adam`, `noisy_sigma_w_adam`, `noisy_mu_b_adam`, `noisy_sigma_b_adam` + - `outcome_w_adam`, `outcome_b_adam` +- PER buffer = `GpuReplayBuffer` (`rl/gpu_replay.rs`); 15 CudaSlice fields incl. `write_head_d, replay_len_d` (both on GPU as `u32`, not host scalars) +- `MultiHorizonLoader`: has `reset(seed)` and `yielded() -> usize`; no `seek(n)` — needs adding for Phase C +- `alpha_rl_train.rs` json!{} block spans ~153 lines starting at `examples/alpha_rl_train.rs:798` + +## Phase split + +| Phase | Scope | Ships | Validation gate | +|-------|-------|-------|-----------------| +| **A** | Eval diag emission | Typed `DiagRecord` + eval loop emits to `eval_diag.jsonl` | Local smoke: train + eval each produce JSONL of expected length, identical schema | +| **B** | Lite checkpoint for inference | `save_inference_checkpoint`, `load_inference_checkpoint`, `IntegratedTrainer::new_for_inference`, smoke binary `alpha_rl_inference_smoke` | Local: save during train → load in separate process → forward pass produces sane output | +| **C** | Full checkpoint + resume | Adam + ISV + scalars + replay (opt-in) + MultiHorizonLoader cursor preserve | Local: split 200-step run into 100+100 via resume; trajectory continues within ε of unbroken run | + +**Each phase is independently shippable.** Phase A is highest priority (the v11 cluster run completed with $61.5k eval pnl but max_dd -$444k and we have ZERO per-step visibility into eval — Phase A would tell us when/why). Phase B is high priority for inference deployment. Phase C is nice-to-have for crash recovery. + +--- + +## Universal rules (apply to all phases) + +1. **All `save_*` methods MUST call `self.stream.synchronize()` before any DtoH download.** This eliminates torn-read corruption when saves happen mid-step. +2. **All `load_*` methods MUST validate length on every uploaded vector** (panic message specifies which buffer mismatched). +3. **Bincode is the serialization format.** No new deps. +4. **No `unwrap()` in new code.** Pre-commit hook will warn; treat as error. +5. **Schema versioning:** every checkpoint struct gets a `version: u32` field. Loaders hard-fail on mismatch. +6. **SHA gating:** loaders log a warning on SHA mismatch but proceed; only refuse on deserialization failure. (Compatibility across SHA is opt-in trust.) + +--- + +## Phase A — Eval-phase diag emission + +### Files + +| Path | Action | +|------|--------| +| `crates/ml-alpha/src/trainer/diag_record.rs` | Create | +| `crates/ml-alpha/src/trainer/mod.rs` | Modify — add `pub mod diag_record;` | +| `crates/ml-alpha/src/trainer/integrated.rs` | Modify — add `snapshot_diag_record` method | +| `crates/ml-alpha/examples/alpha_rl_train.rs` | Modify — replace inline json!{} with typed call; wire eval-diag writer | +| `crates/ml-alpha/tests/eval_diag_emission.rs` | Create | + +### A.1 — Read current json!{} schema (~15 min, no code change) + +- [ ] **Step 1: Dump the existing schema for reference** + +```bash +sed -n '798,1100p' crates/ml-alpha/examples/alpha_rl_train.rs > /tmp/json_block_v11.rs +# Extract every key to a list +grep -oE '"[a-z_]+":' /tmp/json_block_v11.rs | sort -u > /tmp/diag_keys_v11.txt +wc -l /tmp/diag_keys_v11.txt +``` + +This is the source of truth for the DiagRecord schema. The implementer reads `/tmp/diag_keys_v11.txt` while writing the struct in A.2. + +### A.2 — Define typed `DiagRecord` + +- [ ] **Step 1: Add `pub mod diag_record;` to mod.rs** + +```rust +// crates/ml-alpha/src/trainer/mod.rs +pub mod diag_record; +``` + +- [ ] **Step 2: Create the typed struct** + +```rust +// crates/ml-alpha/src/trainer/diag_record.rs +//! Typed per-step diagnostic record. Schema is the same for train and eval. +//! Mirrors the json!{} block in alpha_rl_train.rs:798-950 (v11 baseline). +//! Serializing this struct produces JSONL output that is bit-identical to +//! the previous inline construction modulo serde field-ordering (which +//! consumers ignore). + +use serde::Serialize; + +#[derive(Serialize)] +pub struct DiagRecord { + pub step: u64, + pub elapsed_s: f32, + pub loss: LossBlock, + pub lambdas: LambdasBlock, + pub isv_out: IsvOutBlock, + pub isv_lr: IsvLrBlock, + pub isv_ema_in: IsvEmaInBlock, + pub isv_config: IsvConfigBlock, + pub popart: PopartBlock, + pub ppo: PpoBlock, + pub streaming: StreamingBlock, + pub rewards: RewardsBlock, + pub risk_stack: RiskStackBlock, + pub q_pi_agree_ema: f32, + pub action_hist: Vec, + pub action_entropy: f32, + pub anti_martingale: AntiMartingaleBlock, + pub confidence_gate: ConfidenceGateBlock, + pub controller_branch: ControllerBranchBlock, + pub done_count: u32, + pub frd: FrdBlock, + pub frd_gate: FrdGateBlock, + pub grad_norm_ema: GradNormEmaBlock, + pub k_loop: KLoopBlock, + pub lr_plateau: LrPlateauBlock, + pub outcome_aux: OutcomeAuxBlock, + pub partial_flat: PartialFlatBlock, + pub per_branch_lr: PerBranchLrBlock, + pub position: PositionBlock, + pub position_heat: PositionHeatBlock, + pub pyramid: PyramidBlock, + pub q_bias: QBiasBlock, + pub replay_len: u32, + pub spectral: SpectralBlock, + pub trading: TradingBlock, + pub trail: TrailBlock, + pub units: UnitsBlock, +} + +// ─── Inner blocks ──────────────────────────────────────────────────── +// +// IMPLEMENTER: each block below mirrors the corresponding sub-object in +// the json!{} block. Read alpha_rl_train.rs:798-950 once and write each +// block out. Field types: f32 unless the JSONL shows an integer (use +// u32 or u64). For nullable fields (the v11 diag has a few — e.g. +// risk_stack.regime.transition.remaining), use Option. +// +// All blocks must derive Serialize. Field names must be snake_case and +// match the JSONL keys exactly. + +#[derive(Serialize)] +pub struct LossBlock { + pub bce: f32, + pub q: f32, + pub pi: f32, + pub v: f32, + pub aux: f32, + pub frd: f32, + pub total: f32, +} + +// ... [implementer fills in remaining ~30 inner blocks following the +// same pattern. See /tmp/json_block_v11.rs for the source layout.] +``` + +> **Implementer rule for A.2 Step 2:** every field in the JSONL output must have a matching struct field. After writing, run the schema-comparison check in A.4 Step 4 — that's the gate. + +- [ ] **Step 3: cargo check** + +```bash +SQLX_OFFLINE=true cargo check -p ml-alpha +``` + +Expected: clean. + +- [ ] **Step 4: Commit** + +```bash +git add crates/ml-alpha/src/trainer/diag_record.rs crates/ml-alpha/src/trainer/mod.rs +git commit -m "feat(trainer): typed DiagRecord schema (A.2)" +``` + +### A.3 — Add `IntegratedTrainer::snapshot_diag_record` + +- [ ] **Step 1: Locate the json!{} body and identify every value source** + +```bash +sed -n '798,950p' crates/ml-alpha/examples/alpha_rl_train.rs > /tmp/json_block_v11.rs +``` + +For each field in the json! block, identify where the value comes from: +- Direct trainer field (e.g. `self.last_q_loss`) +- ISV slot read (e.g. `self.isv_mapped.read_record(POPART_SIGMA_INDEX)`) +- DiagStaging readback (the trainer has a `DiagStaging` field that exposes the previous-step's snapshot via `read_*` methods) +- Computed inline in the loop (e.g. `step`, `elapsed_s`) + +- [ ] **Step 2: Implement snapshot method** + +```rust +// crates/ml-alpha/src/trainer/integrated.rs — in impl IntegratedTrainer + +use crate::trainer::diag_record::{DiagRecord, /* every inner block */}; + +impl IntegratedTrainer { + /// Snapshot the trainer state into a typed DiagRecord. Mirrors the + /// json!{} block in alpha_rl_train.rs:798. Used by both train loop + /// (replacing inline construction) and eval loop (NEW — Phase A). + /// + /// IMPORTANT (per universal rule 1): callers running mid-step + /// should expect the values to reflect the PREVIOUS step's data + /// for any field sourced from `DiagStaging` (which uses a separate + /// stream with one-step latency). This is consistent with the + /// existing inline json!{} behavior. + pub fn snapshot_diag_record( + &self, + step: u64, + elapsed_s: f32, + ) -> Result { + Ok(DiagRecord { + step, + elapsed_s, + loss: LossBlock { + bce: self.last_bce_loss, // VERIFY THIS FIELD EXISTS + q: self.last_q_loss, + pi: self.last_pi_loss, + v: self.last_v_loss, + aux: self.last_aux_loss, // VERIFY + frd: self.last_frd_loss, // VERIFY + total: self.last_total_loss, // VERIFY + }, + // ... [implementer fills in by reading the v11 json! block] + }) + } +} +``` + +> **Implementer rule:** before writing each `LossBlock` field, grep for `last_bce_loss`, `last_aux_loss`, `last_frd_loss`, `last_total_loss` in `integrated.rs`. If absent (v1 plan incorrectly assumed all exist — only `last_q_loss`, `last_pi_loss`, `last_v_loss` were verified), the field must be sourced elsewhere. Likely candidates: from the inline json!{} block in alpha_rl_train.rs itself (which already constructs these from somewhere — trace those values back to their source and route through snapshot_diag_record instead). + +- [ ] **Step 3: Build** + +```bash +SQLX_OFFLINE=true cargo check -p ml-alpha +``` + +- [ ] **Step 4: Commit** + +```bash +git add crates/ml-alpha/src/trainer/integrated.rs +git commit -m "feat(trainer): snapshot_diag_record method (A.3)" +``` + +### A.4 — Replace inline json!{} in train loop + wire eval loop + +- [ ] **Step 1: Save baseline diag from v11 cluster run for schema comparison** + +```bash +# Use the existing local smoke output (/tmp/rl-smoke-lpi-diag/diag.jsonl) +head -1 /tmp/rl-smoke-lpi-diag/diag.jsonl | jq -r 'paths(scalars) | join(".")' | sort > /tmp/diag_keys_baseline.txt +wc -l /tmp/diag_keys_baseline.txt +``` + +This is the canonical schema. Any A.4 changes must produce the same set of leaf paths. + +- [ ] **Step 2: Replace train-loop json!{} construction** + +In `alpha_rl_train.rs` around line 798: + +```rust +// BEFORE (existing): +let record = json!({ + "step": step, + "elapsed_s": t_start.elapsed().as_secs_f32(), + // ... 150 more lines +}); +diag_staging.write_record(&record)?; + +// AFTER: +let record = trainer.snapshot_diag_record(step as u64, t_start.elapsed().as_secs_f32())?; +diag_staging.write_record_typed(&record)?; +``` + +Add `write_record_typed` to `DiagStaging`: + +```rust +// crates/ml-alpha/src/trainer/diag_staging.rs + +use crate::trainer::diag_record::DiagRecord; + +impl DiagStaging { + pub fn write_record_typed(&mut self, record: &DiagRecord) -> Result<()> { + let line = serde_json::to_string(record)?; + writeln!(self.writer, "{}", line)?; + Ok(()) + } +} +``` + +- [ ] **Step 3: Wire eval-loop diag writer** + +In `alpha_rl_train.rs` around line 1339 (eval phase entry): + +```rust +let eval_diag_path = cli.eval_diag_jsonl.clone() + .unwrap_or_else(|| cli.out.join("eval_diag.jsonl")); +let mut eval_diag_writer = std::io::BufWriter::new( + std::fs::File::create(&eval_diag_path) + .with_context(|| format!("create {}", eval_diag_path.display()))? +); +eprintln!("eval per-step diag JSONL: {}", eval_diag_path.display()); +``` + +Add the per-step write inside the eval loop: + +```rust +for eval_step in 0..cli.n_eval_steps { + let stats = trainer + .step_with_lobsim_gpu(&mut eval_gpu_loader, &eval_gpu_dataset, &mut sim) + .with_context(|| format!("step_with_lobsim_gpu eval step {eval_step}"))?; + if !stats.l_total.is_finite() { + eprintln!("G8 NaN ABORT during eval at step {eval_step}"); + process::exit(2); + } + // NEW (A.4): emit per-step eval diag. + let record = trainer.snapshot_diag_record(eval_step as u64, t_start.elapsed().as_secs_f32())?; + let line = serde_json::to_string(&record)?; + use std::io::Write; + writeln!(&mut eval_diag_writer, "{}", line)?; + if eval_step % cli.log_every == 0 || eval_step == cli.n_eval_steps - 1 { + eval_diag_writer.flush()?; + eprintln!( + "eval {:>5}/{}: l_total={:.4} elapsed={:.1}s", + eval_step, cli.n_eval_steps, stats.l_total, + t_start.elapsed().as_secs_f32() + ); + } +} +// Flush at end. +eval_diag_writer.flush()?; +drop(eval_diag_writer); +``` + +Note: `step_with_lobsim_gpu` may not populate every field that `snapshot_diag_record` reads (e.g. fields written by the train-phase reward pipeline). Either: +- (a) Verify `snapshot_diag_record` reads only state set by both train and eval step paths +- (b) Have eval emit a smaller `EvalDiagRecord` subset + +> **Implementer:** start with (a). Run the A.4 Step 6 smoke; if eval_diag.jsonl contains stale train-phase values for any field, switch to (b) — extract `EvalDiagRecord` as a subset struct with the same field names but only the fields eval can actually populate. + +- [ ] **Step 4: Add the CLI flag** + +```rust +// crates/ml-alpha/examples/alpha_rl_train.rs — in Cli struct + +/// Eval-phase per-step diag JSONL path. Default: `/eval_diag.jsonl`. +#[arg(long)] +eval_diag_jsonl: Option, +``` + +- [ ] **Step 5: Build** + +```bash +SQLX_OFFLINE=true cargo build --release --example alpha_rl_train -p ml-alpha +``` + +- [ ] **Step 6: Schema-comparison smoke** + +```bash +# Run 100 train + 50 eval +SQLX_OFFLINE=true target/release/examples/alpha_rl_train \ + --n-steps 100 --n-eval-steps 50 --fold-idx 1 --n-folds 3 \ + --mbp10-data-dir /tmp/rl-smoke-lpi-diag/data \ + --predecoded-dir /tmp/rl-smoke-lpi-diag/data \ + --out /tmp/eval-diag-test \ + --instrument-mode front-month --n-backtests 16 --log-every 50 --seed 42 + +# Compare schemas +head -1 /tmp/eval-diag-test/diag.jsonl | jq -r 'paths(scalars) | join(".")' | sort > /tmp/new_train_keys.txt +head -1 /tmp/eval-diag-test/eval_diag.jsonl | jq -r 'paths(scalars) | join(".")' | sort > /tmp/new_eval_keys.txt +diff /tmp/diag_keys_baseline.txt /tmp/new_train_keys.txt # MUST be empty +diff /tmp/new_train_keys.txt /tmp/new_eval_keys.txt # MUST be empty +wc -l /tmp/eval-diag-test/diag.jsonl /tmp/eval-diag-test/eval_diag.jsonl +# Expected: diag.jsonl = 100, eval_diag.jsonl = 50 +``` + +If either diff is non-empty, fix the schema before commit. + +- [ ] **Step 7: Commit** + +```bash +git add crates/ml-alpha/examples/alpha_rl_train.rs crates/ml-alpha/src/trainer/diag_staging.rs +git commit -m "feat(rl): emit eval-phase per-step diag to eval_diag.jsonl (A.4)" +``` + +### A.5 — Integration test + +- [ ] **Step 1: Write smoke test** + +```rust +// crates/ml-alpha/tests/eval_diag_emission.rs + +use anyhow::Result; +use std::process::Command; + +#[test] +#[ignore = "requires CUDA + test_data + alpha_rl_train binary"] +fn eval_diag_emits_per_step_lines_with_train_compatible_schema() -> Result<()> { + let bin = std::env::current_exe()? + .parent().unwrap() // tests/ + .parent().unwrap() // release/ + .join("examples/alpha_rl_train"); + if !bin.exists() { + eprintln!("skipping: alpha_rl_train binary not built ({})", bin.display()); + return Ok(()); + } + + let outdir = tempfile::tempdir()?; + let status = Command::new(&bin) + .args([ + "--n-steps", "100", + "--n-eval-steps", "50", + "--fold-idx", "1", + "--n-folds", "3", + "--mbp10-data-dir", "/tmp/rl-smoke-lpi-diag/data", + "--predecoded-dir", "/tmp/rl-smoke-lpi-diag/data", + "--out", outdir.path().to_str().unwrap(), + "--instrument-mode", "front-month", + "--n-backtests", "16", + "--seed", "42", + "--log-every", "50", + ]) + .env("SQLX_OFFLINE", "true") + .status()?; + assert!(status.success(), "alpha_rl_train failed"); + + let train_diag = outdir.path().join("diag.jsonl"); + let eval_diag = outdir.path().join("eval_diag.jsonl"); + assert!(train_diag.exists(), "diag.jsonl missing"); + assert!(eval_diag.exists(), "eval_diag.jsonl missing"); + + let train_lines = std::fs::read_to_string(&train_diag)?.lines().count(); + let eval_lines = std::fs::read_to_string(&eval_diag)?.lines().count(); + assert_eq!(train_lines, 100, "expected 100 train diag lines, got {}", train_lines); + assert_eq!(eval_lines, 50, "expected 50 eval diag lines, got {}", eval_lines); + + // Schema check: every leaf path in eval line 0 also exists in train line 0. + let train_v: serde_json::Value = serde_json::from_str( + std::fs::read_to_string(&train_diag)?.lines().next().unwrap() + )?; + let eval_v: serde_json::Value = serde_json::from_str( + std::fs::read_to_string(&eval_diag)?.lines().next().unwrap() + )?; + fn collect_paths(v: &serde_json::Value, prefix: &str, out: &mut Vec) { + match v { + serde_json::Value::Object(m) => { + for (k, vv) in m { + let p = if prefix.is_empty() { k.clone() } else { format!("{}.{}", prefix, k) }; + collect_paths(vv, &p, out); + } + } + _ => out.push(prefix.to_string()), + } + } + let mut tp = vec![]; collect_paths(&train_v, "", &mut tp); tp.sort(); + let mut ep = vec![]; collect_paths(&eval_v, "", &mut ep); ep.sort(); + assert_eq!(tp, ep, "train and eval schemas differ"); + + eprintln!("Phase A OK — eval_diag.jsonl emits 50 lines with train-compatible schema"); + Ok(()) +} +``` + +- [ ] **Step 2: Run + commit** + +```bash +SQLX_OFFLINE=true cargo build --release --example alpha_rl_train -p ml-alpha +SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline \ + cargo test --release -p ml-alpha --test eval_diag_emission -- --ignored --nocapture + +git add crates/ml-alpha/tests/eval_diag_emission.rs +git commit -m "test(rl): eval_diag emission smoke (A.5)" +``` + +### A.6 — Push for cluster validation + +- [ ] **Step 1: Push** + +```bash +git push +``` + +- [ ] **Step 2: Submit cluster run** + +Same config as alpha-rl-8ll7j but with eval_diag enabled (default). + +- [ ] **Step 3: Verify on cluster** + +After completion: `kubectl exec` into a fresh investigator pod, verify `eval_diag.jsonl` has 5000 lines on the held-out fold. Sample some lines, ensure they show eval-phase trajectory (mean_pnl decreasing during drawdown, recency resetting, etc.). + +### Phase A done means + +- Schema-comparison diff is empty between train and eval +- 100/50/5000 step counts produce correctly-sized JSONL files +- Cluster run produces `eval_diag.jsonl` alongside `diag.jsonl` +- Can answer "when did the eval max_dd happen?" by querying the new file + +--- + +## Phase B — Lite checkpoint for inference + +### Files + +| Path | Action | +|------|--------| +| `crates/ml-alpha/src/trainer/checkpoint.rs` | Create | +| `crates/ml-alpha/src/trainer/mod.rs` | Modify — add `pub mod checkpoint;` | +| `crates/ml-alpha/src/cfc/trunk.rs` | Modify — add `checkpoint_bytes()` + `load_from_bytes(&mut self, bytes)` | +| `crates/ml-alpha/src/trainer/perception.rs` | Modify — add `trunk_checkpoint_bytes()` + `load_trunk_from_bytes()` | +| `crates/ml-alpha/src/trainer/integrated.rs` | Modify — add `save_inference_checkpoint`, `load_inference_checkpoint`, `new_for_inference` | +| `crates/ml-alpha/build.rs` | Modify — emit `GIT_SHA` env var | +| `crates/ml-alpha/Cargo.toml` | Modify — add `chrono` if missing | +| `crates/ml-alpha/examples/alpha_rl_train.rs` | Modify — `--inference-ckpt-every`, `--checkpoint-keep` flags, save calls in train loop | +| `crates/ml-alpha/examples/alpha_rl_inference_smoke.rs` | Create — load + forward smoke | +| `crates/ml-alpha/tests/checkpoint_lite_roundtrip.rs` | Create | + +### B.1 — Schema (lite only) + +- [ ] **Step 1: Create `checkpoint.rs` with lite-only structs** + +```rust +// crates/ml-alpha/src/trainer/checkpoint.rs + +use serde::{Deserialize, Serialize}; + +pub const CHECKPOINT_VERSION: u32 = 1; + +#[derive(Serialize, Deserialize)] +pub struct LiteCheckpoint { + pub version: u32, + pub sha: String, + pub step: u64, + pub saved_at: String, + pub encoder_bytes: Vec, // CfcTrunk::checkpoint_bytes + pub heads: HeadWeightsCheckpoint, +} + +/// Per-head weight buffers. Field names mirror the trainer's CudaSlice +/// fields. Lengths are NOT stored (validated at load via slice.len()). +#[derive(Serialize, Deserialize, Default)] +pub struct HeadWeightsCheckpoint { + // DqnHead + pub dqn_w: Vec, pub dqn_b: Vec, + pub dqn_w_target: Vec, pub dqn_b_target: Vec, + // IqnHead + pub iqn_w_embed: Vec, pub iqn_b_embed: Vec, + pub iqn_w_out: Vec, pub iqn_b_out: Vec, + pub iqn_w_embed_target: Vec, pub iqn_b_embed_target: Vec, + pub iqn_w_out_target: Vec, pub iqn_b_out_target: Vec, + // DuelingQHead (verified at dueling_q.rs:127-140) + pub dueling_w_v: Vec, pub dueling_b_v: Vec, + pub dueling_w_a: Vec, pub dueling_b_a: Vec, + pub dueling_w_v_target: Vec, pub dueling_b_v_target: Vec, + pub dueling_w_a_target: Vec, pub dueling_b_a_target: Vec, + // PolicyHead + pub policy_w: Vec, pub policy_b: Vec, + // ValueHead + pub value_w: Vec, pub value_b: Vec, +} +``` + +- [ ] **Step 2: Add module to mod.rs** + +```rust +// crates/ml-alpha/src/trainer/mod.rs +pub mod checkpoint; +``` + +- [ ] **Step 3: cargo check + commit** + +```bash +SQLX_OFFLINE=true cargo check -p ml-alpha +git add crates/ml-alpha/src/trainer/checkpoint.rs crates/ml-alpha/src/trainer/mod.rs +git commit -m "feat(trainer): LiteCheckpoint schema (B.1)" +``` + +### B.2 — Build-time SHA + chrono + +- [ ] **Step 1: build.rs emit GIT_SHA + chrono dep** + +```rust +// crates/ml-alpha/build.rs — near the top of main() +let sha = std::process::Command::new("git") + .args(["rev-parse", "--short", "HEAD"]) + .output() + .ok() + .and_then(|o| String::from_utf8(o.stdout).ok()) + .map(|s| s.trim().to_string()) + .unwrap_or_else(|| "unknown".to_string()); +println!("cargo:rustc-env=GIT_SHA={}", sha); +println!("cargo:rerun-if-changed=.git/HEAD"); +// Also watch the actual branch ref so commits trigger rebuild: +println!("cargo:rerun-if-changed=.git/refs/heads"); +``` + +```toml +# crates/ml-alpha/Cargo.toml — verify chrono is present; if not: +chrono = { version = "0.4", features = ["serde"] } +``` + +- [ ] **Step 2: Build to verify GIT_SHA propagates** + +```bash +SQLX_OFFLINE=true cargo check -p ml-alpha +``` + +- [ ] **Step 3: Commit** + +```bash +git add crates/ml-alpha/build.rs crates/ml-alpha/Cargo.toml +git commit -m "build(ml-alpha): emit GIT_SHA env var + chrono dep (B.2)" +``` + +### B.3 — CfcTrunk refactor for in-place load + +Goal: factor `CfcTrunk::save_checkpoint`'s body into `checkpoint_bytes()` (returns `Vec`) and add a new `load_from_bytes(&mut self, bytes: &[u8])` method that uploads into existing buffers (so `IntegratedTrainer::load_inference_checkpoint` can restore without re-allocating GPU memory). + +- [ ] **Step 1: Refactor `save_checkpoint`** + +```rust +// crates/ml-alpha/src/cfc/trunk.rs + +impl CfcTrunk { + /// Build the internal Checkpoint struct + bincode-serialize. + /// Used both by save_checkpoint (writes to file) and by the + /// IntegratedTrainer's embedding path. + pub fn checkpoint_bytes(&self) -> Result> { + // Move existing body of save_checkpoint here, replacing the + // tail + // let bytes = bincode::serialize(&ckpt)?; + // let mut f = ...; + // f.write_all(&bytes)?; + // with + // Ok(bincode::serialize(&ckpt)?) + unimplemented!("move save_checkpoint body here, return bytes") + } + + pub fn save_checkpoint(&self, path: &std::path::Path) -> Result<()> { + use std::io::Write; + let bytes = self.checkpoint_bytes()?; + let mut f = std::fs::File::create(path) + .with_context(|| format!("create {}", path.display()))?; + f.write_all(&bytes).context("write checkpoint bytes")?; + Ok(()) + } +} +``` + +- [ ] **Step 2: Add `load_from_bytes`** + +```rust +impl CfcTrunk { + /// In-place load (no re-allocation). Mirrors `load_checkpoint`'s + /// body but uploads to EXISTING CudaSlice buffers via + /// stream.memcpy_htod. Length-validates every upload. + pub fn load_from_bytes(&mut self, bytes: &[u8]) -> Result<()> { + let ckpt: Checkpoint = bincode::deserialize(bytes) + .context("bincode deserialize CfcTrunk Checkpoint")?; + anyhow::ensure!(ckpt.n_in == self.cfg.n_in, "n_in mismatch"); + anyhow::ensure!(ckpt.n_hid == self.cfg.n_hid, "n_hid mismatch"); + // ... validate every other config field + + // Upload every weight buffer. Use a helper macro: + macro_rules! up { + ($host:expr, $dst:expr) => {{ + anyhow::ensure!($host.len() == $dst.len(), + "{} length mismatch", stringify!($dst)); + self.stream.memcpy_htod(&$host, &mut $dst)?; + }}; + } + up!(ckpt.vsn_w, self.vsn_w_d); + up!(ckpt.vsn_b, self.vsn_b_d); + // ... full enumeration of every field in Checkpoint, mirroring + // the download list in save_checkpoint + Ok(()) + } +} +``` + +- [ ] **Step 3: Build to verify (existing CfcTrunk tests must still pass)** + +```bash +SQLX_OFFLINE=true cargo test -p ml-alpha --lib cfc::trunk -- --nocapture +``` + +- [ ] **Step 4: Commit** + +```bash +git add crates/ml-alpha/src/cfc/trunk.rs +git commit -m "refactor(cfc): factor checkpoint_bytes + add load_from_bytes (B.3)" +``` + +### B.4 — PerceptionTrainer pass-through + +- [ ] **Step 1: Add pass-through methods** + +```rust +// crates/ml-alpha/src/trainer/perception.rs + +impl PerceptionTrainer { + pub fn trunk_checkpoint_bytes(&self) -> Result> { + self.trunk.checkpoint_bytes() + } + pub fn load_trunk_from_bytes(&mut self, bytes: &[u8]) -> Result<()> { + self.trunk.load_from_bytes(bytes) + } +} +``` + +- [ ] **Step 2: Build + commit** + +```bash +SQLX_OFFLINE=true cargo check -p ml-alpha +git add crates/ml-alpha/src/trainer/perception.rs +git commit -m "feat(perception): trunk_checkpoint_bytes/load pass-through (B.4)" +``` + +### B.5 — `save_inference_checkpoint` + `load_inference_checkpoint` + +- [ ] **Step 1: Add download helper + collect_head_weights + save method** + +```rust +// crates/ml-alpha/src/trainer/integrated.rs + +use crate::trainer::checkpoint::{ + LiteCheckpoint, HeadWeightsCheckpoint, CHECKPOINT_VERSION, +}; + +impl IntegratedTrainer { + /// Download a CudaSlice into a host Vec. Synchronizes the + /// trainer stream first to guarantee read-after-write ordering. + fn download_f32(&self, slice: &CudaSlice) -> Result> { + let mut v = vec![0.0_f32; slice.len()]; + self.stream.memcpy_dtoh(slice, v.as_mut_slice())?; + Ok(v) + } + + fn collect_head_weights(&self) -> Result { + Ok(HeadWeightsCheckpoint { + dqn_w: self.download_f32(&self.dqn_head.w_d)?, + dqn_b: self.download_f32(&self.dqn_head.b_d)?, + dqn_w_target: self.download_f32(&self.dqn_head.w_target_d)?, + dqn_b_target: self.download_f32(&self.dqn_head.b_target_d)?, + iqn_w_embed: self.download_f32(&self.iqn_head.w_embed_d)?, + iqn_b_embed: self.download_f32(&self.iqn_head.b_embed_d)?, + iqn_w_out: self.download_f32(&self.iqn_head.w_out_d)?, + iqn_b_out: self.download_f32(&self.iqn_head.b_out_d)?, + iqn_w_embed_target: self.download_f32(&self.iqn_head.w_embed_target_d)?, + iqn_b_embed_target: self.download_f32(&self.iqn_head.b_embed_target_d)?, + iqn_w_out_target: self.download_f32(&self.iqn_head.w_out_target_d)?, + iqn_b_out_target: self.download_f32(&self.iqn_head.b_out_target_d)?, + dueling_w_v: self.download_f32(&self.dueling_q_head.w_v_d)?, + dueling_b_v: self.download_f32(&self.dueling_q_head.b_v_d)?, + dueling_w_a: self.download_f32(&self.dueling_q_head.w_a_d)?, + dueling_b_a: self.download_f32(&self.dueling_q_head.b_a_d)?, + dueling_w_v_target: self.download_f32(&self.dueling_q_head.w_v_target_d)?, + dueling_b_v_target: self.download_f32(&self.dueling_q_head.b_v_target_d)?, + dueling_w_a_target: self.download_f32(&self.dueling_q_head.w_a_target_d)?, + dueling_b_a_target: self.download_f32(&self.dueling_q_head.b_a_target_d)?, + policy_w: self.download_f32(&self.policy_head.w_d)?, + policy_b: self.download_f32(&self.policy_head.b_d)?, + value_w: self.download_f32(&self.value_head.w_d)?, + value_b: self.download_f32(&self.value_head.b_d)?, + }) + } + + /// Save weights-only checkpoint for inference deployment. + /// MUST synchronize the train stream first (universal rule 1) to + /// avoid torn reads when called mid-step. + pub fn save_inference_checkpoint( + &self, + path: &std::path::Path, + step: u64, + ) -> Result<()> { + use std::io::Write; + self.stream.synchronize() + .context("sync train stream before checkpoint download")?; + let encoder_bytes = self.perception.trunk_checkpoint_bytes()?; + let heads = self.collect_head_weights()?; + let ckpt = LiteCheckpoint { + version: CHECKPOINT_VERSION, + sha: env!("GIT_SHA").to_string(), + step, + saved_at: chrono::Utc::now().to_rfc3339(), + encoder_bytes, + heads, + }; + let bytes = bincode::serialize(&ckpt)?; + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent)?; + } + let mut f = std::fs::File::create(path) + .with_context(|| format!("create {}", path.display()))?; + f.write_all(&bytes)?; + Ok(()) + } +} +``` + +- [ ] **Step 2: Add upload helper + apply_head_weights + load method** + +```rust +impl IntegratedTrainer { + fn upload_f32(stream: &CudaStream, host: &[f32], dst: &mut CudaSlice) -> Result<()> { + anyhow::ensure!(host.len() == dst.len(), + "upload_f32 length mismatch: got {} expected {}", host.len(), dst.len()); + stream.memcpy_htod(host, dst)?; + Ok(()) + } + + fn apply_head_weights(&mut self, h: &HeadWeightsCheckpoint) -> Result<()> { + let s = &self.stream; + Self::upload_f32(s, &h.dqn_w, &mut self.dqn_head.w_d)?; + Self::upload_f32(s, &h.dqn_b, &mut self.dqn_head.b_d)?; + Self::upload_f32(s, &h.dqn_w_target, &mut self.dqn_head.w_target_d)?; + Self::upload_f32(s, &h.dqn_b_target, &mut self.dqn_head.b_target_d)?; + Self::upload_f32(s, &h.iqn_w_embed, &mut self.iqn_head.w_embed_d)?; + Self::upload_f32(s, &h.iqn_b_embed, &mut self.iqn_head.b_embed_d)?; + Self::upload_f32(s, &h.iqn_w_out, &mut self.iqn_head.w_out_d)?; + Self::upload_f32(s, &h.iqn_b_out, &mut self.iqn_head.b_out_d)?; + Self::upload_f32(s, &h.iqn_w_embed_target, &mut self.iqn_head.w_embed_target_d)?; + Self::upload_f32(s, &h.iqn_b_embed_target, &mut self.iqn_head.b_embed_target_d)?; + Self::upload_f32(s, &h.iqn_w_out_target, &mut self.iqn_head.w_out_target_d)?; + Self::upload_f32(s, &h.iqn_b_out_target, &mut self.iqn_head.b_out_target_d)?; + Self::upload_f32(s, &h.dueling_w_v, &mut self.dueling_q_head.w_v_d)?; + Self::upload_f32(s, &h.dueling_b_v, &mut self.dueling_q_head.b_v_d)?; + Self::upload_f32(s, &h.dueling_w_a, &mut self.dueling_q_head.w_a_d)?; + Self::upload_f32(s, &h.dueling_b_a, &mut self.dueling_q_head.b_a_d)?; + Self::upload_f32(s, &h.dueling_w_v_target, &mut self.dueling_q_head.w_v_target_d)?; + Self::upload_f32(s, &h.dueling_b_v_target, &mut self.dueling_q_head.b_v_target_d)?; + Self::upload_f32(s, &h.dueling_w_a_target, &mut self.dueling_q_head.w_a_target_d)?; + Self::upload_f32(s, &h.dueling_b_a_target, &mut self.dueling_q_head.b_a_target_d)?; + Self::upload_f32(s, &h.policy_w, &mut self.policy_head.w_d)?; + Self::upload_f32(s, &h.policy_b, &mut self.policy_head.b_d)?; + Self::upload_f32(s, &h.value_w, &mut self.value_head.w_d)?; + Self::upload_f32(s, &h.value_b, &mut self.value_head.b_d)?; + Ok(()) + } + + /// Load a lite checkpoint. Returns the saved train step. + /// Hard-fails on version mismatch; logs (does not fail) on SHA mismatch. + pub fn load_inference_checkpoint(&mut self, path: &std::path::Path) -> Result { + let bytes = std::fs::read(path) + .with_context(|| format!("read {}", path.display()))?; + let ckpt: LiteCheckpoint = bincode::deserialize(&bytes) + .context("bincode deserialize LiteCheckpoint")?; + anyhow::ensure!( + ckpt.version == CHECKPOINT_VERSION, + "checkpoint version mismatch: got {}, expected {}", + ckpt.version, CHECKPOINT_VERSION + ); + if ckpt.sha != env!("GIT_SHA") { + eprintln!("checkpoint SHA mismatch (saved={} current={}) — proceeding", + ckpt.sha, env!("GIT_SHA")); + } + self.perception.load_trunk_from_bytes(&ckpt.encoder_bytes)?; + self.apply_head_weights(&ckpt.heads)?; + self.stream.synchronize()?; + Ok(ckpt.step) + } +} +``` + +- [ ] **Step 3: Build** + +```bash +SQLX_OFFLINE=true cargo check -p ml-alpha +``` + +- [ ] **Step 4: Commit** + +```bash +git add crates/ml-alpha/src/trainer/integrated.rs +git commit -m "feat(trainer): save/load_inference_checkpoint (B.5)" +``` + +### B.6 — Round-trip unit test + +- [ ] **Step 1: Test** + +```rust +// crates/ml-alpha/tests/checkpoint_lite_roundtrip.rs + +use anyhow::Result; +use ml_alpha::trainer::integrated::IntegratedTrainer; +use ml_core::device::MlDevice; +use tempfile::tempdir; + +#[test] +#[ignore = "requires CUDA"] +fn lite_checkpoint_save_load_preserves_weights() -> Result<()> { + let Ok(dev) = MlDevice::cuda(0) else { return Ok(()); }; + + // Trainer A — seed 42. + let mut a = IntegratedTrainer::with_controllers_bootstrapped(&dev, 16, 42)?; + let a_dqn_w = download_for_test(&a, &a.dqn_head.w_d)?; + let a_policy_w = download_for_test(&a, &a.policy_head.w_d)?; + let a_dueling_w_v_tgt = download_for_test(&a, &a.dueling_q_head.w_v_target_d)?; + + // Save. + let tmp = tempdir()?; + let path = tmp.path().join("test.lite.ckpt"); + a.save_inference_checkpoint(&path, 1234)?; + + // Trainer B — different seed. + let mut b = IntegratedTrainer::with_controllers_bootstrapped(&dev, 16, 999)?; + let b_dqn_w_pre = download_for_test(&b, &b.dqn_head.w_d)?; + assert_ne!(a_dqn_w, b_dqn_w_pre, "different seeds → different initial weights"); + + // Load A into B. + let loaded = b.load_inference_checkpoint(&path)?; + assert_eq!(loaded, 1234); + + // Verify all three sampled buffers now match. + let b_dqn_w_post = download_for_test(&b, &b.dqn_head.w_d)?; + let b_policy_w_post = download_for_test(&b, &b.policy_head.w_d)?; + let b_dueling_w_v_tgt_post = download_for_test(&b, &b.dueling_q_head.w_v_target_d)?; + for (i, (x, y)) in a_dqn_w.iter().zip(&b_dqn_w_post).enumerate() { + assert!((x - y).abs() < 1e-6, "dqn_w[{}] differs: {} vs {}", i, x, y); + } + for (i, (x, y)) in a_policy_w.iter().zip(&b_policy_w_post).enumerate() { + assert!((x - y).abs() < 1e-6, "policy_w[{}] differs", i); + } + for (i, (x, y)) in a_dueling_w_v_tgt.iter().zip(&b_dueling_w_v_tgt_post).enumerate() { + assert!((x - y).abs() < 1e-6, "dueling target[{}] differs", i); + } + + eprintln!("Phase B OK — lite checkpoint round-trips weights bit-exact"); + Ok(()) +} + +// Public download helper for tests. Lives in integrated.rs as +// `pub fn download_f32_for_test` (under `#[cfg(any(test, feature = "test-utils"))]` +// — must NOT be `#[cfg(test)]` because integration tests are a separate +// crate context). +fn download_for_test( + trainer: &IntegratedTrainer, + slice: &cudarc::driver::CudaSlice, +) -> Result> { + trainer.download_f32_for_test(slice) +} +``` + +In `integrated.rs`, expose the download helper for tests: + +```rust +impl IntegratedTrainer { + /// Test-only download helper. Gated by `test-utils` feature so it + /// compiles for integration tests but NOT in production binaries. + #[cfg(any(test, feature = "test-utils"))] + pub fn download_f32_for_test(&self, slice: &CudaSlice) -> Result> { + self.download_f32(slice) + } +} +``` + +In `Cargo.toml`: + +```toml +[features] +test-utils = [] + +[dev-dependencies] +ml-alpha = { path = ".", features = ["test-utils"] } +# (or enable via [dependencies] when running tests — simpler: +# uncomment the `pub fn download_f32_for_test` cfg gate and use +# `[[test]] required-features = ["test-utils"]` if needed.) +``` + +> **Implementer:** the cleanest path is to enable `test-utils` for integration tests via `[[test]] ... required-features = ["test-utils"]` per test, OR drop the cfg and make `download_f32_for_test` permanently `pub` (with a doc comment "intended for tests only"). Cargo's interaction between integration tests and crate features is fiddly; pick whichever path your local build accepts. + +- [ ] **Step 2: Run + commit** + +```bash +SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline \ + cargo test --release -p ml-alpha --test checkpoint_lite_roundtrip -- --ignored --nocapture +git add crates/ml-alpha/tests/checkpoint_lite_roundtrip.rs crates/ml-alpha/src/trainer/integrated.rs crates/ml-alpha/Cargo.toml +git commit -m "test(trainer): lite checkpoint round-trip (B.6)" +``` + +### B.7 — Wire CLI flags + train-loop saves + +- [ ] **Step 1: CLI flags** + +```rust +// crates/ml-alpha/examples/alpha_rl_train.rs — in Cli struct + +/// Save a lite (inference) checkpoint every N train steps. 0 disables. +#[arg(long, default_value_t = 5000)] +inference_ckpt_every: usize, + +/// Rolling-window count of lite checkpoints to keep on disk. +#[arg(long, default_value_t = 3)] +checkpoint_keep: usize, +``` + +- [ ] **Step 2: Save calls in train loop** + +```rust +// In the train loop, after each step: + +if cli.inference_ckpt_every > 0 && step > 0 && step % cli.inference_ckpt_every == 0 { + let ckpt_dir = cli.out.join("checkpoints"); + let path = ckpt_dir.join(format!("inference_step_{:06}.lite.ckpt", step)); + trainer.save_inference_checkpoint(&path, step as u64)?; + eprintln!("saved lite checkpoint {} ({} MB)", + path.display(), + std::fs::metadata(&path)?.len() / (1024 * 1024)); + rolling_cleanup(&ckpt_dir, "inference_step_", ".lite.ckpt", cli.checkpoint_keep)?; +} +``` + +- [ ] **Step 3: End-of-phase saves** + +```rust +// After train loop, before eval (or after if no eval): +let final_path = cli.out.join("checkpoints/inference_final.lite.ckpt"); +trainer.save_inference_checkpoint(&final_path, cli.n_steps as u64)?; +eprintln!("inference_final saved"); + +// After eval loop: +let post_eval_path = cli.out.join("checkpoints/inference_post_eval.lite.ckpt"); +trainer.save_inference_checkpoint(&post_eval_path, cli.n_steps as u64)?; +eprintln!("inference_post_eval saved"); +``` + +- [ ] **Step 4: rolling_cleanup helper** + +```rust +// crates/ml-alpha/examples/alpha_rl_train.rs — top-level fn + +fn rolling_cleanup(dir: &std::path::Path, prefix: &str, suffix: &str, keep: usize) -> Result<()> { + if !dir.exists() { return Ok(()); } + let mut entries: Vec<_> = std::fs::read_dir(dir)? + .filter_map(|e| e.ok()) + .filter(|e| { + let name = e.file_name().to_string_lossy().into_owned(); + name.starts_with(prefix) && name.ends_with(suffix) + }) + .collect(); + entries.sort_by_key(|e| e.metadata().and_then(|m| m.modified()).ok()); + while entries.len() > keep { + let oldest = entries.remove(0); + let _ = std::fs::remove_file(oldest.path()); + } + Ok(()) +} +``` + +- [ ] **Step 5: Build + local smoke** + +```bash +SQLX_OFFLINE=true cargo build --release --example alpha_rl_train -p ml-alpha + +# 200 steps with --inference-ckpt-every 100 +SQLX_OFFLINE=true target/release/examples/alpha_rl_train \ + --n-steps 200 --inference-ckpt-every 100 --checkpoint-keep 2 \ + --mbp10-data-dir /tmp/rl-smoke-lpi-diag/data \ + --predecoded-dir /tmp/rl-smoke-lpi-diag/data \ + --out /tmp/lite-ckpt-test --instrument-mode front-month \ + --n-backtests 16 --log-every 50 --seed 42 + +ls -lh /tmp/lite-ckpt-test/checkpoints/ +# Expected: ~3 files: inference_step_000100, inference_step_000200, inference_final +``` + +- [ ] **Step 6: Commit** + +```bash +git add crates/ml-alpha/examples/alpha_rl_train.rs +git commit -m "feat(rl): inference-ckpt-every flag + train-loop saves (B.7)" +``` + +### B.8 — Inference smoke binary + +- [ ] **Step 1: Create smoke binary** + +```rust +// crates/ml-alpha/examples/alpha_rl_inference_smoke.rs + +//! Load a lite checkpoint and run N inference steps on held-out data. + +use anyhow::{Context, Result}; +use clap::Parser; +use ml_alpha::trainer::integrated::IntegratedTrainer; +use ml_alpha::data::loader::{discover_mbp10_files_sorted, MultiHorizonLoader, MultiHorizonLoaderConfig}; +use ml_alpha::data::gpu_dataset::{GpuDataset, GpuDataLoader}; +use ml_core::device::MlDevice; +use std::path::PathBuf; + +#[derive(Parser)] +struct Cli { + #[arg(long)] checkpoint: PathBuf, + #[arg(long)] mbp10_data_dir: PathBuf, + #[arg(long)] predecoded_dir: PathBuf, + #[arg(long, default_value_t = 1000)] n_steps: usize, + #[arg(long, default_value_t = 16)] n_backtests: usize, + #[arg(long, default_value_t = 32)] seq_len: usize, + #[arg(long, default_value_t = 0)] gpu_idx: usize, +} + +fn main() -> Result<()> { + let cli = Cli::parse(); + let dev = MlDevice::cuda(cli.gpu_idx as i32)?; + eprintln!("CUDA initialised on device {}", cli.gpu_idx); + + let mut trainer = IntegratedTrainer::with_controllers_bootstrapped(&dev, cli.n_backtests, 0)?; + let loaded_step = trainer.load_inference_checkpoint(&cli.checkpoint)?; + eprintln!("loaded checkpoint @ train_step={}", loaded_step); + + // Build loader (mirror alpha_rl_train.rs loader setup). + let files = discover_mbp10_files_sorted(&cli.mbp10_data_dir) + .context("discover mbp10 files")?; + let loader = MultiHorizonLoader::new(&MultiHorizonLoaderConfig { + files, + predecoded_dir: cli.predecoded_dir.clone(), + seq_len: cli.seq_len, + n_max_sequences: cli.n_steps, + // ... fill in remaining config from alpha_rl_train.rs + .. Default::default() + })?; + let gpu_dataset = loader.upload_to_gpu(&dev, cli.n_backtests)?; + let mut gpu_loader = GpuDataLoader::new(&dev, cli.n_backtests)?; + let mut sim = /* build LobSimCuda — mirror alpha_rl_train.rs */; + + for step in 0..cli.n_steps { + let stats = trainer + .step_with_lobsim_gpu(&mut gpu_loader, &gpu_dataset, &mut sim) + .with_context(|| format!("step {step}"))?; + if !stats.l_total.is_finite() { + eprintln!("NaN at step {}", step); + std::process::exit(2); + } + if step % 100 == 0 || step == cli.n_steps - 1 { + eprintln!("inference step {:>5}/{}: l_total={:.4}", step, cli.n_steps, stats.l_total); + } + } + eprintln!("inference smoke OK: {} steps with loaded checkpoint", cli.n_steps); + Ok(()) +} +``` + +- [ ] **Step 2: Build + run against B.7's saved checkpoint** + +```bash +SQLX_OFFLINE=true cargo build --release --example alpha_rl_inference_smoke -p ml-alpha + +SQLX_OFFLINE=true target/release/examples/alpha_rl_inference_smoke \ + --checkpoint /tmp/lite-ckpt-test/checkpoints/inference_final.lite.ckpt \ + --mbp10-data-dir /tmp/rl-smoke-lpi-diag/data \ + --predecoded-dir /tmp/rl-smoke-lpi-diag/data \ + --n-steps 200 --n-backtests 16 +``` + +Expected: "inference smoke OK: 200 steps with loaded checkpoint" without abort. + +- [ ] **Step 3: Commit** + +```bash +git add crates/ml-alpha/examples/alpha_rl_inference_smoke.rs +git commit -m "feat(rl): alpha_rl_inference_smoke binary (B.8)" +``` + +### B.9 — Cluster validation + +- [ ] **Step 1: Push** + +```bash +git push +``` + +- [ ] **Step 2: Submit cluster run** + +Use same flags as v11 plus `--inference-ckpt-every 5000`. Verify `argo-alpha-rl.sh` passes the flag through (read the script — if it doesn't, add the pass-through and re-commit). + +- [ ] **Step 3: After completion, verify checkpoint files on cluster + run inference smoke against `inference_post_eval.lite.ckpt`** + +### Phase B done means + +- Round-trip test passes on local RTX 3050 (bit-exact weight restoration) +- Local 200-step run produces 3 lite checkpoints (last 2 + final per rolling cleanup) +- `alpha_rl_inference_smoke` runs 200 inference steps on the saved checkpoint without error +- Cluster run produces `inference_post_eval.lite.ckpt` + 5+ rolling lite checkpoints +- A separate cluster pod can be invoked to load + run the post-eval checkpoint + +--- + +## Phase C — Full checkpoint + resume + +### Scope + +This is the longest phase. Goals: +1. Add full checkpoint variant: lite + Adam state + ISV bus + trainer scalars +2. Add **optional** replay buffer save/load (gated on `--checkpoint-replay`) +3. Add **MultiHorizonLoader cursor preservation** so resume continues at the correct data position +4. Add `--resume-from` CLI flag +5. Validate resume continuity within ε tolerance + +### Files (additive on top of Phase B) + +| Path | Action | +|------|--------| +| `crates/ml-alpha/src/trainer/checkpoint.rs` | Modify — add `FullCheckpoint`, `AdamStateCheckpoint`, `ReplayBufferCheckpoint`, `TrainerScalarsCheckpoint` | +| `crates/ml-alpha/src/trainer/optim.rs` | Modify — add `AdamW::save_state`, `AdamW::load_state` | +| `crates/ml-alpha/src/trainer/perception.rs` | Modify — add `collect_adam_states`, `apply_adam_state` | +| `crates/ml-alpha/src/trainer/integrated.rs` | Modify — add `save_full_checkpoint`, `load_full_checkpoint`, `collect_adam_state`, `apply_adam_state`, `collect_replay_buffer`, `apply_replay_buffer` | +| `crates/ml-alpha/src/rl/gpu_replay.rs` | Modify — add `snapshot_for_checkpoint`/`restore_from_checkpoint` | +| `crates/ml-alpha/src/data/loader.rs` | Modify — add `MultiHorizonLoader::cursor() -> usize` and `with_cursor(cfg, cursor)` | +| `crates/ml-alpha/examples/alpha_rl_train.rs` | Modify — `--checkpoint-every`, `--checkpoint-replay`, `--resume-from` flags; resume entry path; full-checkpoint save calls | +| `crates/ml-alpha/tests/checkpoint_full_roundtrip.rs` | Create | +| `crates/ml-alpha/tests/checkpoint_replay_roundtrip.rs` | Create | +| `crates/ml-alpha/tests/resume_continuation.rs` | Create | + +### Tasks + +C.1 — `FullCheckpoint` schema (additive on `checkpoint.rs`) +C.2 — `AdamW::save_state` / `load_state` (already partially in v1 plan Task 1.2 — copy as-is, it's correct) +C.3 — `IntegratedTrainer::collect_adam_state` / `apply_adam_state` over the 22 Adams enumerated above; PerceptionTrainer pass-through for any perception-side Adams +C.4 — `save_full_checkpoint` / `load_full_checkpoint` (mirrors B.5 but adds Adam + ISV + scalars) +C.5 — `MultiHorizonLoader::cursor()` getter + `with_cursor(cfg, cursor)` constructor +C.6 — `GpuReplayBuffer::snapshot_for_checkpoint` / `restore_from_checkpoint` over its 15 buffers (download/upload pattern; `write_head_d` / `replay_len_d` are CudaSlice, not host scalars) +C.7 — CLI integration: `--checkpoint-every`, `--checkpoint-replay`, `--resume-from`; resume entry; periodic save calls +C.8 — Unit tests: full round-trip (ISV + Adam + scalars), replay round-trip (deterministic sample post-restore), resume continuation +C.9 — Cluster validation: save mid-run, simulate restart, verify ETA matches unbroken + +### Why deferred behind B + +Phase C touches 8 files vs Phase B's 5, includes data-loader internals, and has the trickiest correctness criterion (resume must continue without trajectory divergence). Phase B + Phase A together cover the immediate needs (eval visibility, inference deployment). Phase C is purely about training continuity, which we've never actually needed in production yet. + +Phase C will get its own detailed plan once Phase A + B land and we have real cluster experience with the saved lite checkpoints. + +--- + +## Universal validation gate + +Each phase must pass these before merging into the next: + +1. All new code compiles clean (no warnings about `unwrap()` from pre-commit hook) +2. All new unit tests pass on local RTX 3050 +3. Local smoke produces expected files +4. Cluster run completes producing expected artifacts +5. Pearl saved + MEMORY.md updated documenting the design rationale + +## Done means + +- Phase A: cluster run produces `eval_diag.jsonl` with 5000 lines, train and eval schemas match +- Phase B: separate inference-smoke pod loads `inference_post_eval.lite.ckpt` and runs forward without abort +- Phase C: simulated crash + resume produces a final loss within ε of the unbroken-run baseline diff --git a/docs/superpowers/plans/2026-05-31-checkpoints-and-eval-diag.md b/docs/superpowers/plans/2026-05-31-checkpoints-and-eval-diag.md new file mode 100644 index 000000000..f254228f2 --- /dev/null +++ b/docs/superpowers/plans/2026-05-31-checkpoints-and-eval-diag.md @@ -0,0 +1,1612 @@ +# IntegratedTrainer Checkpoints + Eval Diag — Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Add full save/load surface to `IntegratedTrainer` (lite + full checkpoint variants), wire CLI for save/resume in `alpha_rl_train`, and make eval phase emit per-step diag to `eval_diag.jsonl`. + +**Architecture:** Two checkpoint kinds (lite = weights only for inference; full = lite + Adam moments + ISV bus + scalars + optional replay). Bincode flat-binary, schema versioned. Diag emission extracted to a typed `DiagRecord` struct producible by trainer; train and eval loops each write to their own JSONL. + +**Tech Stack:** Rust, cudarc, bincode, serde, AdamW (existing). + +**Spec:** `docs/superpowers/specs/2026-05-31-checkpoints-and-eval-diag-design.md` + +**Branch:** `ml-alpha-checkpoints-eval-diag` (off `ml-alpha-regime-observer`) + +--- + +## File structure + +| Path | Action | Responsibility | +|------|--------|----------------| +| `crates/ml-alpha/src/trainer/checkpoint.rs` | Create | All `*Checkpoint` structs, `IntegratedCheckpoint` top-level, version constant | +| `crates/ml-alpha/src/trainer/diag_record.rs` | Create | Typed `DiagRecord` struct; serde::Serialize → one JSONL line | +| `crates/ml-alpha/src/trainer/optim.rs` | Modify | Add `AdamW::save_state` and `AdamW::load_state` | +| `crates/ml-alpha/src/trainer/integrated.rs` | Modify | Add `save_inference_checkpoint`, `save_full_checkpoint`, `load_checkpoint`, `snapshot_diag_record` methods | +| `crates/ml-alpha/src/trainer/mod.rs` | Modify | `pub mod checkpoint; pub mod diag_record;` | +| `crates/ml-alpha/examples/alpha_rl_train.rs` | Modify | New CLI flags; resume entry; save cadence; eval-diag writer | +| `crates/ml-alpha/tests/checkpoint_lite_roundtrip.rs` | Create | Phase 1 unit test | +| `crates/ml-alpha/tests/checkpoint_full_roundtrip.rs` | Create | Phase 2 unit test | +| `crates/ml-alpha/tests/checkpoint_replay_roundtrip.rs` | Create | Phase 3 unit test | +| `crates/ml-alpha/tests/eval_diag_emission.rs` | Create | Phase 4 smoke | +| `crates/ml-alpha/tests/resume_continuation.rs` | Create | Phase 5 smoke | + +--- + +## Phase 1: Schema + lite save/load + +### Task 1.1: Define `IntegratedCheckpoint` schema + +**Files:** +- Create: `crates/ml-alpha/src/trainer/checkpoint.rs` +- Modify: `crates/ml-alpha/src/trainer/mod.rs` + +- [ ] **Step 1: Add `pub mod checkpoint;` to trainer/mod.rs** + +```rust +// crates/ml-alpha/src/trainer/mod.rs — append at end of pub mod list +pub mod checkpoint; +``` + +- [ ] **Step 2: Create `checkpoint.rs` with all schema structs** + +```rust +// crates/ml-alpha/src/trainer/checkpoint.rs +//! Bincode-serialized checkpoint format for IntegratedTrainer. +//! +//! Two variants: +//! - "lite": weights only (encoder + heads + targets). For inference. +//! - "full": lite + Adam moments + ISV bus + trainer scalars + optional replay. +//! +//! See docs/superpowers/specs/2026-05-31-checkpoints-and-eval-diag-design.md. + +use serde::{Deserialize, Serialize}; + +/// Format version — bump on any layout-breaking change. +pub const CHECKPOINT_VERSION: u32 = 1; + +#[derive(Serialize, Deserialize)] +pub struct IntegratedCheckpoint { + pub version: u32, + pub sha: String, + pub step: u64, + pub saved_at: String, + pub kind: String, // "lite" or "full" + pub encoder: CfcTrunkCheckpointWire, + pub heads: HeadWeightsCheckpoint, + pub adam: Option, + pub isv: Option>, + pub scalars: Option, + pub replay: Option, +} + +/// Mirrors crate::cfc::trunk::Checkpoint (already bincode-serialized). +/// We embed by re-serializing into bytes so a struct-level change in +/// trunk::Checkpoint doesn't force an IntegratedCheckpoint version bump +/// — the trunk's own loader validates its bytes. Trade-off: one extra +/// serialize/deserialize per save/load (negligible cost). +#[derive(Serialize, Deserialize)] +pub struct CfcTrunkCheckpointWire { + pub bytes: Vec, +} + +#[derive(Serialize, Deserialize, Default)] +pub struct HeadWeightsCheckpoint { + // DqnHead + pub dqn_w: Vec, + pub dqn_b: Vec, + pub dqn_w_target: Vec, + pub dqn_b_target: Vec, + // IqnHead + pub iqn_w_embed: Vec, + pub iqn_b_embed: Vec, + pub iqn_w_out: Vec, + pub iqn_b_out: Vec, + pub iqn_w_embed_target: Vec, + pub iqn_b_embed_target: Vec, + pub iqn_w_out_target: Vec, + pub iqn_b_out_target: Vec, + // DuelingQHead + pub dueling_w_v: Vec, + pub dueling_b_v: Vec, + pub dueling_w_a: Vec, + pub dueling_b_a: Vec, + pub dueling_w_v_target: Vec, + pub dueling_b_v_target: Vec, + pub dueling_w_a_target: Vec, + pub dueling_b_a_target: Vec, + // PolicyHead + pub policy_w: Vec, + pub policy_b: Vec, + // ValueHead + pub value_w: Vec, + pub value_b: Vec, +} + +#[derive(Serialize, Deserialize, Default)] +pub struct AdamStateCheckpoint { + /// One entry per optimizer slot the trainer holds; key names mirror + /// trainer fields (e.g. "dqn_w", "iqn_w_embed", ...). Each value is + /// (m_d, v_d, step_count). + pub states: Vec, +} + +#[derive(Serialize, Deserialize)] +pub struct NamedAdamState { + pub name: String, + pub m: Vec, + pub v: Vec, + pub step: i32, +} + +#[derive(Serialize, Deserialize, Default)] +pub struct TrainerScalarsCheckpoint { + pub last_q_loss: f32, + pub last_pi_loss: f32, + pub last_v_loss: f32, + pub train_step: u64, +} + +#[derive(Serialize, Deserialize, Default)] +pub struct ReplayBufferCheckpoint { + /// PER priority sumtree (length 2 × capacity − 1). + pub priority_tree: Vec, + /// Per-sample feature vectors flattened (capacity × feature_dim). + pub features: Vec, + /// Per-sample next-state features. + pub next_features: Vec, + /// Per-sample taken action indices. + pub actions: Vec, + /// Per-sample rewards. + pub rewards: Vec, + /// Per-sample done flags (0.0 or 1.0 stored as f32 for buffer + /// uniformity). + pub dones: Vec, + /// Current write index in the circular buffer. + pub write_head: u64, + /// Number of samples actually populated (≤ capacity). + pub n_populated: u64, +} +``` + +- [ ] **Step 3: Run `SQLX_OFFLINE=true cargo check -p ml-alpha` to verify schema compiles** + +Expected: `Finished dev [unoptimized] target(s) in