diff --git a/docs/superpowers/specs/2026-05-27-alpha-rl-perf-checkpoint-walkforward.md b/docs/superpowers/specs/2026-05-27-alpha-rl-perf-checkpoint-walkforward.md new file mode 100644 index 000000000..4a0c1cdf6 --- /dev/null +++ b/docs/superpowers/specs/2026-05-27-alpha-rl-perf-checkpoint-walkforward.md @@ -0,0 +1,152 @@ +# Alpha-RL Performance + Checkpoint + Walk-Forward + +**Date**: 2026-05-27 +**Status**: Approved +**Scope**: TF32 matmuls, checkpoint/resume, controller fusion, walk-forward validation + +## Context + +The alpha-rl pipeline achieves wr=0.567 at b=1024 but takes 4.5h for 100k steps on L40S at 6.2 sps — and the model converges by 5-10k steps. The last 90k steps are wasted. The L40S node OOM'd at 58k from monitoring pod overhead, losing all state (zero checkpoint infrastructure). All cuBLAS SGEMMs force `CUBLAS_COMPUTE_32F`, bypassing Tensor Cores entirely on both L40S and H100. + +## Goals + +1. **2-3× throughput** via TF32 Tensor Core acceleration on cuBLAS SGEMMs +2. **Crash recovery** via checkpoint/resume every 5k steps +3. **~30% launch overhead reduction** via controller kernel fusion +4. **OOS validation** via 3-fold walk-forward at 25k steps/fold on H100 + +Target wall clock: 3 folds × 25k steps at ~15 sps on H100 = ~83 min total (vs 4.5h+ single-fold L40S that never finished). + +## P0: TF32 Matmuls + +### What + +Enable TF32 Tensor Core math on all cuBLAS handles. TF32 uses 19-bit mantissa (vs FP32's 23-bit) with identical range. Precision loss is well within RL gradient noise. PyTorch default since 1.7. + +### Where + +Two cuBLAS handle creation sites: + +1. `crates/ml-alpha/src/rl/dqn.rs` — after `cublasCreate_v2()`, add `cublasSetMathMode(handle, CUBLAS_TF32_TENSOR_OP_MATH)`. Affects DQN forward (online + target) and backward (grad_h + grad_w) = 4 SGEMMs/step. + +2. `crates/ml-alpha/src/mamba2_block.rs` — same pattern after handle creation. Affects W_in, W_a, W_b, W_out projection GEMMs = 4+ SGEMMs/step. + +### Impact + +8+ SGEMMs per step go from FP32 scalar to TF32 Tensor Core. On L40S (Ada Lovelace): ~2× throughput. On H100 (Hopper): ~3× throughput. Combined with H100's 2× memory bandwidth: expect 12-20 sps at b=1024 (vs 6.2 current). + +### Verification + +Run 1k-step local smoke on RTX 3050 Ti (sm_86, supports TF32). Compare l_q at step 1000 with and without TF32 — should be within 1%. + +## P1: Checkpoint/Resume + +### What + +Serialize training state every 5k steps to PVC. Resume from last checkpoint on restart. + +### State to persist + +| Component | Size | Source | +|-----------|------|--------| +| DQN weights (online) | W[128×231] + b[231] = ~120KB | `dqn_head.w_d`, `dqn_head.b_d` | +| DQN weights (target) | Same = ~120KB | `dqn_head.w_target_d`, `dqn_head.b_target_d` | +| Policy head weights | ~120KB | `policy_head.w_d`, `policy_head.b_d` | +| V head weights | ~2KB | `v_head.w_d`, `v_head.b_d` | +| Encoder weights | ~2-10MB (Mamba2 + CfC) | `encoder.params_d()` | +| Adam state (m, v per param) | 2× model size = ~20MB | Per-group m_d, v_d | +| ISV bus | 585 × f32 = 2.3KB | `isv_dev_ptr` | +| PER buffer | priorities + tree = ~1MB | `gpu_replay.priorities_d` | +| Step counter + RNG | ~100B | `step`, xorshift state | + +Total: ~50MB per checkpoint. + +### Format + +Flat binary: `[magic: u32][version: u32][step: u64][n_sections: u32][sections...]` where each section is `[name_len: u16][name: bytes][data_len: u64][data: bytes]`. No serde, no JSON — raw device→mapped-pinned→file. + +### Path + +`/feature-cache/alpha-rl-runs//fold/checkpoint-.bin` + +Keep last 2 checkpoints (rolling). Delete older ones to save PVC space. + +### Resume + +CLI flag `--resume-from `. Loads checkpoint, restores all device buffers, continues from saved step. The `alpha_rl_train.rs` binary checks for the flag before the training loop. + +### Verification + +Save at step 1000, kill, resume, compare l_q/wr at step 2000 with an uninterrupted run. Should be identical (deterministic with scoped_init_seed per pearl). + +## P2: Controller Kernel Fusion + +### What + +Fuse 10+ single-thread ISV controller kernels into one `rl_fused_all_controllers` mega-kernel. Currently each controller is a separate `raw_launch(grid=1, block=1)` — the launch overhead (~5μs each) dominates the actual compute (~1μs each). + +### Controllers to fuse + +- `rl_gamma_controller` +- `rl_target_tau_controller` +- `rl_ppo_clip_controller` (legacy but still launched) +- `rl_entropy_coef_controller` +- `rl_per_alpha_controller` +- `rl_reward_scale_controller` +- `rl_rollout_steps_controller` +- `rl_lr_controller` (5 heads) + +All share the same pattern: read ISV input slot, Wiener-α blend, Schulman bounded step, write ISV output slot. The existing `rl_fused_controllers.cu` already handles a subset — extend to cover all. + +### Impact + +~10 kernel launches → 1. Saves ~50μs/step. At 6 sps that's ~0.3ms saved per 160ms step = ~0.2% improvement. Small but free. + +### Verification + +Before/after nsys trace: total GPU kernel launches should drop by ~10 per step. + +## P3: Walk-Forward Validation + +### What + +3-fold temporal walk-forward on H100 with 25k steps per fold. Each fold trains on window [i] and evaluates on window [i+1]. The argo template already supports `--folds 3` which fans out via DAG. + +### Config + +```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 + # --folds 3 via argo-train.sh +``` + +### Success criteria + +- wr > 0.55 on ALL 3 folds (not just the average) +- No fold with wr < 0.50 +- Entropy stable (no collapse on any fold) +- hold% between 30-70% on all folds + +### Wall clock estimate + +3 folds × 25k steps. With TF32 on H100 at ~15 sps: 25k/15 = 28 min/fold. Sequential: 84 min. Parallel (3 H100s): 28 min. + +## Implementation Order + +1. **P0: TF32** — 2 one-line changes, local smoke, deploy +2. **P2: Controller fusion** — extend existing fused kernel, local smoke +3. **P1: Checkpoint** — new infrastructure, needs careful testing +4. **P3: Walk-forward** — deploy after P0+P2 are validated + +P0 and P2 can be implemented in parallel. P1 is independent. P3 runs after all code changes are merged. + +## Anti-patterns to avoid + +- No FP16 matmuls (loss scaling complexity, RL numerical sensitivity) +- No multi-GPU (NCCL overhead not justified at 25k steps) +- No changes to model architecture (validated at wr=0.567) +- No monitoring pods on GPU node (caused OOM — use log-only monitoring)