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>
6.0 KiB
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
- 2-3× throughput via TF32 Tensor Core acceleration on cuBLAS SGEMMs
- Crash recovery via checkpoint/resume every 5k steps
- ~30% launch overhead reduction via controller kernel fusion
- 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:
-
crates/ml-alpha/src/rl/dqn.rs— aftercublasCreate_v2(), addcublasSetMathMode(handle, CUBLAS_TF32_TENSOR_OP_MATH). Affects DQN forward (online + target) and backward (grad_h + grad_w) = 4 SGEMMs/step. -
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_controllerrl_target_tau_controllerrl_ppo_clip_controller(legacy but still launched)rl_entropy_coef_controllerrl_per_alpha_controllerrl_reward_scale_controllerrl_rollout_steps_controllerrl_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
- P0: TF32 — 2 one-line changes, local smoke, deploy
- P2: Controller fusion — extend existing fused kernel, local smoke
- P1: Checkpoint — new infrastructure, needs careful testing
- 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)