# 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%