Files
foxhunt/docs/superpowers/specs/2026-05-27-alpha-rl-perf-checkpoint-walkforward.md
jgrusewski 6e641b934c spec: alpha-rl perf + checkpoint + walk-forward design
P0: TF32 matmuls (2-3× SGEMM throughput, 1 line each in dqn.rs + mamba2_block.rs)
P1: Checkpoint/resume every 5k steps (~50MB flat binary to PVC)
P2: Controller kernel fusion (10 launches → 1)
P3: Walk-forward 3×25k on H100 (~84 min total)

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2026-05-27 20:10:47 +02:00

6.0 KiB
Raw Blame History

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/<sha>/fold<N>/checkpoint-<step>.bin

Keep last 2 checkpoints (rolling). Delete older ones to save PVC space.

Resume

CLI flag --resume-from <path>. 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

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)