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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -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 `<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**
|
||||
|
||||
```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<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`:
|
||||
```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 <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%
|
||||
Reference in New Issue
Block a user