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>
This commit is contained in:
jgrusewski
@@ -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/<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
|
||||||
|
|
||||||
|
```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)
|
||||||
Reference in New Issue
Block a user