From e730b5cbd197a4efb3665064c7a7b274155e29d8 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Wed, 27 May 2026 20:16:48 +0200 Subject: [PATCH] plan: alpha-rl perf + checkpoint + walk-forward implementation 7 tasks: TF32 on DQN + Mamba2 handles, AdamW save/load, trainer checkpoint/resume, training loop wiring, local smoke, H100 walk-forward. Co-Authored-By: Claude Opus 4.7 --- ...27-alpha-rl-perf-checkpoint-walkforward.md | 333 ++++++++++++++++++ 1 file changed, 333 insertions(+) create mode 100644 docs/superpowers/plans/2026-05-27-alpha-rl-perf-checkpoint-walkforward.md diff --git a/docs/superpowers/plans/2026-05-27-alpha-rl-perf-checkpoint-walkforward.md b/docs/superpowers/plans/2026-05-27-alpha-rl-perf-checkpoint-walkforward.md new file mode 100644 index 000000000..58b7f6e2c --- /dev/null +++ b/docs/superpowers/plans/2026-05-27-alpha-rl-perf-checkpoint-walkforward.md @@ -0,0 +1,333 @@ +# Alpha-RL Performance + Checkpoint + Walk-Forward 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:** 2-3× training throughput via TF32, crash recovery via checkpointing, OOS validation via 3-fold walk-forward on H100. + +**Architecture:** Enable TF32 Tensor Core math on all cuBLAS handles (2 sites). Add checkpoint/resume serialization for model weights + Adam state + ISV to the training binary. Validate with 3-fold walk-forward at 25k steps/fold on H100. + +**Tech Stack:** Rust, CUDA cuBLAS, cudarc, mapped-pinned memory, Argo Workflows + +--- + +### Task 1: Enable TF32 on DQN cuBLAS handle + +**Files:** +- Modify: `crates/ml-alpha/src/rl/dqn.rs:314` (after `cublasSetWorkspace_v2`) + +- [ ] **Step 1: Add TF32 math mode after workspace setup** + +In `crates/ml-alpha/src/rl/dqn.rs`, after line 313 (the closing of the `cublasSetWorkspace_v2` block), add inside the same `unsafe` block: + +```rust + cudarc::cublas::sys::cublasSetMathMode( + *cublas.handle(), + cudarc::cublas::sys::cublasMath_t::CUBLAS_TF32_TENSOR_OP_MATH, + ) + .result() + .map_err(|e| anyhow::anyhow!("DqnHead: cublasSetMathMode TF32: {e:?}"))?; +``` + +- [ ] **Step 2: Build check** + +Run: `SQLX_OFFLINE=true cargo check -p ml-alpha` +Expected: compiles cleanly + +- [ ] **Step 3: Commit** + +```bash +git add crates/ml-alpha/src/rl/dqn.rs +git commit -m "perf(cuda): enable TF32 Tensor Core math on DQN cuBLAS handle" +``` + +### Task 2: Enable TF32 on Mamba2 cuBLAS handle + +**Files:** +- Modify: `crates/ml-alpha/src/mamba2_block.rs:549` (after `cublasSetWorkspace_v2`) + +- [ ] **Step 1: Add TF32 math mode after workspace setup** + +In `crates/ml-alpha/src/mamba2_block.rs`, after line 549 (the closing of the `cublasSetWorkspace_v2` block), add inside the same `unsafe` block: + +```rust + cudarc::cublas::sys::cublasSetMathMode( + *cublas.handle(), + cudarc::cublas::sys::cublasMath_t::CUBLAS_TF32_TENSOR_OP_MATH, + ) + .result() + .map_err(|e| anyhow!("Mamba2Block: cublasSetMathMode TF32: {e:?}"))?; +``` + +- [ ] **Step 2: Build check** + +Run: `SQLX_OFFLINE=true cargo check -p ml-alpha` +Expected: compiles cleanly + +- [ ] **Step 3: Local smoke test (RTX 3050 Ti, sm_86 supports TF32)** + +Run a 100-step local smoke and verify l_q is within 1% of baseline: +```bash +FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml-alpha --lib -- smoke_tests --ignored --nocapture 2>&1 | tail -5 +``` +Expected: smoke passes, no NaN + +- [ ] **Step 4: Commit** + +```bash +git add crates/ml-alpha/src/mamba2_block.rs +git commit -m "perf(cuda): enable TF32 Tensor Core math on Mamba2 cuBLAS handle" +``` + +### Task 3: Add AdamW save/load methods + +**Files:** +- Modify: `crates/ml-alpha/src/trainer/optim.rs` + +- [ ] **Step 1: Add save method to AdamW** + +Add after the existing `impl AdamW` block (or at the end of the impl): + +```rust + /// Serialize Adam state (m, v, step_count, hyperparams) to writer. + /// Reads device buffers via mapped-pinned DtoH (one-shot, not hot path). + pub fn save(&self, w: &mut impl std::io::Write) -> Result<()> { + let m_host = self._stream.clone().read_sync(&self.m) + .map_err(|e| anyhow::anyhow!("AdamW save m: {e}"))?; + let v_host = self._stream.clone().read_sync(&self.v) + .map_err(|e| anyhow::anyhow!("AdamW save v: {e}"))?; + let n = m_host.len() as u64; + w.write_all(&n.to_le_bytes())?; + w.write_all(&self.step_count_host.to_le_bytes())?; + w.write_all(&self.lr.to_le_bytes())?; + w.write_all(&self.beta1.to_le_bytes())?; + w.write_all(&self.beta2.to_le_bytes())?; + w.write_all(&self.eps.to_le_bytes())?; + w.write_all(&self.wd.to_le_bytes())?; + let m_bytes: &[u8] = bytemuck::cast_slice(&m_host); + w.write_all(m_bytes)?; + let v_bytes: &[u8] = bytemuck::cast_slice(&v_host); + w.write_all(v_bytes)?; + Ok(()) + } + + /// Restore Adam state from reader. Sizes must match. + pub fn load(&mut self, r: &mut impl std::io::Read) -> Result<()> { + let mut buf8 = [0u8; 8]; + let mut buf4 = [0u8; 4]; + r.read_exact(&mut buf8)?; + let n = u64::from_le_bytes(buf8) as usize; + anyhow::ensure!(n == self.m.len(), "AdamW load: m size mismatch {n} vs {}", self.m.len()); + r.read_exact(&mut buf4)?; + self.step_count_host = i32::from_le_bytes(buf4); + r.read_exact(&mut buf4)?; self.lr = f32::from_le_bytes(buf4); + r.read_exact(&mut buf4)?; self.beta1 = f32::from_le_bytes(buf4); + r.read_exact(&mut buf4)?; self.beta2 = f32::from_le_bytes(buf4); + r.read_exact(&mut buf4)?; self.eps = f32::from_le_bytes(buf4); + r.read_exact(&mut buf4)?; self.wd = f32::from_le_bytes(buf4); + let mut m_host = vec![0f32; n]; + r.read_exact(bytemuck::cast_slice_mut(&mut m_host))?; + self._stream.clone().write_sync(&m_host, &mut self.m) + .map_err(|e| anyhow::anyhow!("AdamW load m: {e}"))?; + let mut v_host = vec![0f32; n]; + r.read_exact(bytemuck::cast_slice_mut(&mut v_host))?; + self._stream.clone().write_sync(&v_host, &mut self.v) + .map_err(|e| anyhow::anyhow!("AdamW load v: {e}"))?; + Ok(()) + } +``` + +- [ ] **Step 2: Verify bytemuck is in dependencies** + +Run: `grep bytemuck crates/ml-alpha/Cargo.toml` +If missing, add `bytemuck = { version = "1", features = ["derive"] }`. + +- [ ] **Step 3: Build check** + +Run: `SQLX_OFFLINE=true cargo check -p ml-alpha` + +- [ ] **Step 4: Commit** + +```bash +git add crates/ml-alpha/src/trainer/optim.rs crates/ml-alpha/Cargo.toml +git commit -m "feat(rl): AdamW save/load for checkpoint persistence" +``` + +### Task 4: Add IntegratedTrainer checkpoint save/load + +**Files:** +- Modify: `crates/ml-alpha/src/trainer/integrated.rs` + +- [ ] **Step 1: Add save_checkpoint method** + +Add a `save_checkpoint` method to `IntegratedTrainer` that serializes: +- DQN weights (online + target): `dqn_head.w_d`, `dqn_head.b_d`, `dqn_head.w_target_d`, `dqn_head.b_target_d` +- Policy/V head weights +- ISV bus (585 f32 from mapped-pinned `isv_mapped.host_ptr`) +- All 20 AdamW states via `adam.save()` +- Step counter +- Encoder via `perception.save_checkpoint()` + +Format: sequential sections with `[name_len: u16][name: bytes][data_len: u64][data: bytes]`. +File header: `[magic: u32 = 0x464F5843][version: u32 = 1][step: u64]`. + +Write to `/checkpoint-.bin`. Keep last 2, delete older. + +- [ ] **Step 2: Add load_checkpoint method** + +Inverse of save: read header, validate magic/version, restore each section by name lookup. Skip unknown sections for forward compatibility. + +- [ ] **Step 3: Build check** + +Run: `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(rl): IntegratedTrainer checkpoint save/load" +``` + +### Task 5: Wire checkpoint into training loop + +**Files:** +- Modify: `crates/ml-alpha/examples/alpha_rl_train.rs` + +- [ ] **Step 1: Add CLI flags** + +Add to the CLI struct (clap): +```rust + /// Resume from checkpoint file path. Restores model weights, Adam + /// state, ISV, and step counter. Training continues from saved step. + #[arg(long)] + resume_from: Option, + + /// Checkpoint interval in steps (default: 5000). Set to 0 to disable. + #[arg(long, default_value = "5000")] + checkpoint_every: usize, +``` + +- [ ] **Step 2: Add resume logic before training loop** + +After trainer init, before `for step in 0..cli.n_steps`: +```rust + let start_step = if let Some(ref ckpt_path) = cli.resume_from { + trainer.load_checkpoint(ckpt_path) + .with_context(|| format!("resume from {}", ckpt_path.display()))? + } else { + 0 + }; +``` + +Change loop to `for step in start_step..cli.n_steps`. + +- [ ] **Step 3: Add checkpoint save inside training loop** + +After the per-step JSONL write, before the next iteration: +```rust + if cli.checkpoint_every > 0 && step > 0 && step % cli.checkpoint_every == 0 { + let ckpt_path = cli.out.join(format!("checkpoint-{step}.bin")); + trainer.save_checkpoint(&ckpt_path, step as u64) + .with_context(|| format!("checkpoint at step {step}"))?; + eprintln!("checkpoint saved: {}", ckpt_path.display()); + // Rolling: keep last 2, delete older + let old_step = step.saturating_sub(cli.checkpoint_every * 2); + if old_step > 0 { + let old_path = cli.out.join(format!("checkpoint-{old_step}.bin")); + let _ = std::fs::remove_file(&old_path); + } + } +``` + +- [ ] **Step 4: Build check** + +Run: `SQLX_OFFLINE=true cargo check -p ml-alpha` + +- [ ] **Step 5: Commit** + +```bash +git add crates/ml-alpha/examples/alpha_rl_train.rs +git commit -m "feat(rl): wire checkpoint save/resume into training loop" +``` + +### Task 6: Local validation smoke + +**Files:** None (test only) + +- [ ] **Step 1: Run 200-step local smoke with TF32** + +```bash +SQLX_OFFLINE=true cargo build --release -p ml-alpha --example alpha_rl_train +# Run 200 steps with checkpoint at 100 +./target/release/examples/alpha_rl_train \ + --n-backtests 16 --n-steps 200 --checkpoint-every 100 \ + --out /tmp/smoke-tf32 \ + --mbp10-data-dir test_data/futures-baseline \ + --trades-data-dir test_data/futures-baseline 2>&1 | tail -5 +``` + +Expected: completes without NaN, checkpoint files at `/tmp/smoke-tf32/checkpoint-100.bin` + +- [ ] **Step 2: Test resume** + +```bash +./target/release/examples/alpha_rl_train \ + --n-backtests 16 --n-steps 200 --checkpoint-every 100 \ + --resume-from /tmp/smoke-tf32/checkpoint-100.bin \ + --out /tmp/smoke-tf32-resume \ + --mbp10-data-dir test_data/futures-baseline \ + --trades-data-dir test_data/futures-baseline 2>&1 | tail -5 +``` + +Expected: starts from step 100, runs to 200, no crash + +- [ ] **Step 3: Commit all together** + +```bash +git add -A +git commit -m "test: validate TF32 + checkpoint save/resume local smoke" +``` + +### Task 7: Deploy walk-forward on H100 + +**Files:** None (deployment only) + +- [ ] **Step 1: Push branch** + +```bash +git push origin ml-alpha-phase-a +``` + +- [ ] **Step 2: Submit 3-fold walk-forward on H100** + +```bash +./scripts/argo-train.sh \ + --model alpha-rl \ + --branch ml-alpha-phase-a \ + --gpu-pool ci-training-h100 \ + --folds 3 +``` + +If `argo-train.sh` doesn't support `--folds` with step override, submit directly: +```bash +argo submit --from=wftmpl/alpha-rl -n foxhunt \ + -p git-branch=ml-alpha-phase-a \ + -p n-steps=25000 \ + -p n-backtests=1024 \ + -p per-capacity=65536 \ + -p gpu-pool=ci-training-h100 +``` +Run 3 times with fold-idx 0, 1, 2. + +- [ ] **Step 3: Monitor with log-only (no extra pods on GPU node)** + +```bash +kubectl logs -n foxhunt --tail=1 -f | grep "step.*wr=" +``` + +- [ ] **Step 4: Validate success criteria** + +All 3 folds must show: +- wr > 0.55 +- No fold with wr < 0.50 +- Entropy stable (no collapse) +- hold% between 30-70%