foxhunt/docs/superpowers/plans/2026-05-27-alpha-rl-perf-checkpoint-walkforward.md

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:

            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

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:

            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:

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

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):

    /// 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

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 <out_dir>/checkpoint-<step>.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

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):

    /// Resume from checkpoint file path. Restores model weights, Adam
    /// state, ISV, and step counter. Training continues from saved step.
    #[arg(long)]
    resume_from: Option<PathBuf>,

    /// 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:

    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:

        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

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

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

./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

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

git push origin ml-alpha-phase-a

• Step 2: Submit 3-fold walk-forward on H100

./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:

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)

kubectl logs <pod-name> -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%