diff --git a/docs/superpowers/specs/2026-03-21-h100-epoch-optimization-design.md b/docs/superpowers/specs/2026-03-21-h100-epoch-optimization-design.md new file mode 100644 index 000000000..fdfc109d9 --- /dev/null +++ b/docs/superpowers/specs/2026-03-21-h100-epoch-optimization-design.md @@ -0,0 +1,196 @@ +# H100 Epoch Optimization: 37s → <5s per epoch + +**Date**: 2026-03-21 +**Target**: <5s/epoch on H100 (80GB), <10s acceptable, <3s stretch goal +**Current**: 37.56s/epoch (H100, 51 atoms, 300 training steps, CUDA Graph) + +## Profiling Results + +RTX 3050 phase breakdown (proportions apply to H100): + +| Phase | Time | % | Per-step | +|-------|-----:|--:|---------| +| Init (data upload) | 10ms | 0.1% | once/epoch | +| Experience collection | 348ms | 2.3% | 1 kernel launch | +| **Training steps** | **14,858ms** | **97.2%** | 300 steps | +| Validation | 27ms | 0.2% | 1 backtest pass | + +H100 epoch time: 37.56s × 97.2% training = **36.5s in training steps**. +At 300 steps/epoch: **122ms per training step**. + +## Root Cause Analysis + +### Primary bottleneck: PER sampling CPU roundtrips (300×/epoch) + +Each `sample_proportional()` call does: +1. **DtoH**: `cs_total()` reads prefix sum total (1 scalar) — forces stream sync +2. **CPU RNG**: `rand::thread_rng().gen::()` × batch_size — serial CPU work +3. **HtoD**: uploads random thresholds (batch_size × 4 bytes) +4. **HtoD**: zeros `batch_max` (4 bytes) + +300 calls × ~100μs per sync barrier = **30ms just for sync**, but the CPU RNG + alloc + memcpy overhead adds 100-400ms per call. + +### Secondary: kernel underutilization + +Training kernel: `block_dim=(32,1,1)`, `grid_dim=(batch_size,1,1)` = 256 blocks × 32 threads = **8,192 threads on 132-SM H100** (3.0% occupancy). Each SM can run 2048 threads — we're using 1 warp per SM on a fraction of 132 SMs. + +### Tertiary: sequential phase execution + +Experience collection → Training → Validation runs serially. No overlap. + +## Design: Three-Phase Optimization + +### Phase 1: Eliminate CPU Roundtrips (target: 10-15s/epoch) + +#### 1a. GPU-native PER threshold generation + +Replace CPU `rand::thread_rng()` with Philox CUDA RNG kernel (already have this pattern in IQN tau sampler): + +```cuda +// New kernel: per_generate_thresholds_kernel +__global__ void per_generate_thresholds( + float* __restrict__ thresholds, // [batch_size] output + float total_sum, // from prefix sum (stays on GPU) + unsigned long long seed, // counter-based, no state + int batch_size) +{ + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= batch_size) return; + // Philox-based PRNG (same as iqn_sample_taus_kernel) + unsigned int r = philox_uint(seed, i); + float u = __uint_as_float(0x3f800000u | (r >> 9)) - 1.0f; // uniform [0,1) + thresholds[i] = u * total_sum; +} +``` + +**Key change**: `total_sum` stays on GPU — read it via a single DtoD copy from `cs_buf[n-1]` to a scalar buffer, then pass as kernel arg. Zero DtoH. + +#### 1b. Fuse PER kernels into single launch + +Current: 10+ kernel launches per sample (pow_alpha, prefix_sum, searchsorted, 5 gathers, IS weights, reduce). + +Fuse into 2 kernels: +1. `per_sample_and_gather_kernel`: prefix sum total (already computed) → thresholds → searchsorted → gather all fields +2. `per_is_weights_kernel`: compute IS weights + normalize + +Or capture PER sampling into a second CUDA Graph. + +#### 1c. Remove adam_step HtoD + +Store adam_step as a GPU scalar, increment via 1-element kernel or atomicAdd. + +#### 1d. Audit and remove all hot-path `to_host()` calls + +Remove `GpuTensor::to_host()` from the public API. Replace with `to_host_checkpoint()` (clearly marked for cold-path use only). Compile-time guard via `#[deprecated]` or feature gate. + +### Phase 2: Increase Kernel Occupancy (target: 5-8s/epoch) + +#### 2a. Batched training: multiple samples per warp + +Current: 1 warp (32 threads) per sample, 64 samples per graph replay. +Target: 4-8 samples per block (128-256 threads), or increase batch_size to 512-1024. + +With batch_size=512: `grid=(512,1,1), block=(32,1,1)` = 16K threads = 7.5% occupancy. +With batch_size=1024 + 4 samples/block: 8K threads at 128 threads/block = better SM utilization. + +#### 2b. cuBLAS for network layers + +Replace warp-cooperative matvec with cuBLAS `cublasSgemmStridedBatched` for the shared trunk layers (256×256 matmul). For batch_size=1024, this is a `[1024, 256] × [256, 256]` GEMM — cuBLAS tensor cores will dominate. + +Trade-off: cuBLAS adds launch overhead but uses tensor cores (3.9 PFLOPS BF16 on H100 vs ~100 TFLOPS custom kernel). + +Profile both and pick winner. Custom kernel may win for batch<256. + +#### 2c. BF16 throughout (not just storage) + +Current: BF16 weight storage → F32 accumulation in kernel. +Target: BF16 weights + BF16 activations + F32 accumulation (tensor core matmul). + +H100 tensor cores: BF16 GEMM at 989 TFLOPS vs F32 at 67 TFLOPS = **14.8× potential**. + +### Phase 3: Pipeline Overlap (target: <5s/epoch) + +#### 3a. Dual-stream: overlap experience + training + +``` +Stream A: [experience_epoch_N+1] → [experience_epoch_N+2] → ... +Stream B: [train_epoch_N] → [train_epoch_N+1] → ... + ↑ uses experiences from epoch N (already in buffer) +``` + +After epoch 1, experience collection for epoch N+1 overlaps with training for epoch N. ~2× throughput. + +#### 3b. Fuse PER sampling + training into mega-graph + +Capture the entire sequence (PER sample → batch upload → forward+loss → backward → adam → EMA) into a single CUDA Graph. One `cuGraphLaunch` per step instead of 15+ kernel launches. + +Requires: GPU-native PER (Phase 1a) and graph-compatible memory pool. + +#### 3c. Async validation + +Move validation backtest to a separate stream, overlapping with the start of the next epoch's experience collection. + +## Implementation Order + +1. **Phase 1a**: GPU Philox RNG for PER thresholds (eliminate sync barriers) +2. **Phase 2a**: Increase batch_size to 512+ on H100 (config change, immediate win, amplifies all later optimizations) +3. **Phase 1b**: Fuse PER kernels + pre-allocate all PER buffers (zero `cuMemAlloc` after warmup) +4. **Phase 1c**: GPU adam_step increment +5. **Phase 1d**: Remove `to_host()` from hot paths, deprecate the function +6. **Phase 2b**: Profile cuBLAS vs custom for training matmul +7. **Phase 2c**: BF16 activations (with C51 log_softmax numerical guard) +8. **Phase 3a**: Dual-stream overlap (requires resolving stream unification tension) +9. **Phase 3b**: Mega-graph (if 3a insufficient) + +## Expected Speedup Stack + +| Optimization | Expected speedup | Cumulative | +|-------------|:----------------:|:----------:| +| GPU PER RNG (eliminate 300 sync barriers) | 2-3× | 12-18s | +| Batch_size 512+ (occupancy) | 1.3-1.5× | 9-12s | +| Fuse PER kernels + pre-alloc (reduce 14→2 launches) | 1.2-1.5× | 6-10s | +| cuBLAS/tensor core matmul | 1.5-3× | 3-6s | +| Dual-stream overlap | 1.5-2× | 2-4s | + +Conservative estimate: **5-8s/epoch**. Aggressive: **3-4s/epoch**. + +## Known Risks + +### CUDA Graph capture requires fixed buffer addresses +PER's `sample_proportional` currently allocates new `CudaSlice` buffers per call (lines 434, 436, 443, 450-451, 478, 490 — ~10 allocs × 300 calls = 3,000 GPU malloc/free per epoch). For mega-graph capture (Phase 3b), ALL buffers must be pre-allocated at fixed addresses. Phase 1b must pre-allocate these. + +### Dual-stream conflicts with stream unification +The `stream_unification_plan.md` explicitly advocates collapsing to a single forked stream to avoid cudarc event tracking conflicts. Phase 3a proposes re-introducing dual streams. Resolution: use separate `CudaContext` per stream (each context has independent event tracking), or complete cudarc event tracking fix first. + +### C51 log_softmax underflow in BF16 +Phase 2c (BF16 activations) risks numerical issues in C51 distributional loss. `log_softmax` on BF16 can underflow for low-probability atoms. Mitigation: keep loss computation in F32 accumulation, only use BF16 for matmul operands. Profile numerically before deploying. + +### Prefix sum single-block scalability +The prefix sum kernel runs with `grid_dim=(1,1,1)` — single block processing up to ~1M elements. At large replay buffer capacities (1M+ on H100), this serializes. May need multi-block parallel scan (CUB `DeviceScan`) for Phase 2a batch_size increases. + +## Success Criteria + +- [ ] Phase breakdown log shows zero DtoH in training steps (except 8-byte loss+grad_norm per step) +- [ ] Zero `cuMemAlloc` calls in `sample_proportional` after first warmup call (pre-allocated buffers only) +- [ ] `scripts/gpu-hotpath-guard.sh` catches any new CPU roundtrip in hot path +- [ ] H100 epoch time < 5s with 300 training steps, 51 atoms, batch_size ≥ 512 +- [ ] RTX 3050 epoch time < 10s (proportional improvement from current ~15s) +- [ ] All 1861+ tests pass +- [ ] Training metrics regression < 5%: eval Sharpe within 5% of baseline, loss convergence within 10% of baseline epoch count +- [ ] `to_host()` deprecated with `#[deprecated]` attribute — only `to_host_checkpoint()` available for cold paths + +## Risk: Experience Collector Becomes Bottleneck + +At 2.3% of current epoch time, the experience collector seems negligible. But on H100 with 256 episodes × 500 timesteps = 128K forward passes, it likely takes 5-10s. Once training drops to <5s, experience collection becomes the dominant phase. + +Mitigations (included in Phase 3): +- **Phase 3a (dual-stream)**: overlap experience with training → hides latency entirely +- **Occupancy**: experience kernel uses `episodes_per_block=4-8` with warp kernel on H100 — already reasonably optimized +- **Reduce timesteps**: 500 timesteps may be excessive; 200 could suffice with higher episode count for same total experiences +- **Profile independently**: add per-kernel CUDA event timing inside experience collector to identify sub-kernel bottlenecks + +## Non-Goals + +- Multi-GPU training (separate effort, needs NCCL) +- Model architecture changes (network dims stay 256×256) +- Hyperopt speed (separate binary, different bottleneck)