diff --git a/docs/superpowers/specs/2026-04-01-h100-gpu-optimization-design.md b/docs/superpowers/specs/2026-04-01-h100-gpu-optimization-design.md new file mode 100644 index 000000000..742acf7c1 --- /dev/null +++ b/docs/superpowers/specs/2026-04-01-h100-gpu-optimization-design.md @@ -0,0 +1,151 @@ +# H100 GPU Pipeline Optimization Design + +**Goal:** Reduce DQN training epoch time from 2 min to <10s on H100 by eliminating GPU idle time, maximizing SM utilization (52% → 90%+), and exploiting H100's multi-stream capabilities. + +**Architecture:** Async CUDA event pipeline replaces blocking synchronize calls. Multi-stream dispatch for independent network passes. CUDA Graph capture for experience collection. Larger batch sizes to saturate tensor cores. + +**Current state:** H100 PCIe, 80GB VRAM, 52% SM utilization, 0% memory bandwidth, 102W/350W TDP. 341 training steps/epoch at batch=1024. 204K experiences collected per epoch via 100-timestep host-driven loop. + +--- + +## Root Cause Analysis + +The primary bottleneck is **CPU→GPU serialization**: `stream.synchronize()` is called 341 times per epoch in `replay_adam_and_readback()`. The H100 completes each graph pair in microseconds but then idles while Rust re-enters the loop, re-acquires locks, and reissues kernels. Secondary: the experience collection has a 100-iteration host loop dispatching 4-5 kernels each, creating launch latency bubbles. + +nvidia-smi confirms: 52% SM utilization (should be 90%+), 0% memory bandwidth (no large tensor ops saturating HBM3), 102W/350W TDP (GPU is thermally idling). + +--- + +## Optimizations (ordered by impact) + +### 1. CUDA Events replace per-step stream.synchronize() + +**File:** `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs:2872` + +**Current:** After every training step, `stream.synchronize()` blocks the CPU until both CUDA graphs complete. Loss + grad_norm scalars are read back synchronously. The CPU then re-enters Rust, re-acquires locks, and reissues the next step. This creates a bubble between every step. + +**Fix:** Record a CUDA event after `graph_adam` completes. The CPU immediately proceeds to issue the next step's batch sampling + graph launches without waiting. Loss/grad scalars from step N are read when event N signals completion (typically before step N+2 needs them for PER priority update). + +**Validation:** Loss values must be identical to the synchronous path. Run 10 epochs on RTX 3050 and compare loss trajectories bit-for-bit. + +**Risk:** PER priority updates lag by 1-2 steps. Acceptable — PER priorities are approximate by design (stale priorities are standard in distributed RL). + +### 2. CUDA Graph for experience collection + +**File:** `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs:1506` + +**Current:** `launch_timestep_loop` runs a host-driven `for t in 0..100` loop. Each iteration dispatches 4-5 kernels through Rust. The GPU executes each kernel batch quickly but waits for the CPU to increment `t`, compute pointers, and reissue. + +**Fix:** Capture the entire 100-timestep loop as a single CUDA Graph. All per-timestep parameters (epsilon, step counter, state buffer pointers) become GPU-resident buffers updated via a single DtoH memcpy before graph launch. One graph launch replaces 400-500 individual kernel dispatches. + +**Complexity:** High. The experience kernel uses dynamic epsilon (decays per step) and reads from the replay buffer which changes size. These must be converted to GPU-side indirection (pointer-to-pointer) or parameterized via graph node updates (CUDA 12+ `cudaGraphExecKernelNodeSetParams`). + +**Validation:** Compare collected experiences (states, actions, rewards) with the sequential path for 1 epoch. Must be numerically identical given the same RNG seed. + +**Risk:** Graph capture may fail if any kernel uses dynamic shared memory or conditional branches that differ between timesteps. Fall back to the sequential path if capture fails. + +### 3. Multi-stream branch parallelism in forward/backward pass + +**File:** `crates/ml/src/cuda_pipeline/batched_forward.rs:270` + +**Current:** The 3 advantage branches (exposure=9, order=3, urgency=3) are dispatched sequentially in a `for d in 0..3` loop. Each branch does: GEMM → bias+ReLU → GEMM → bias (4 ops). Total: 12 sequential ops. + +**Fix:** Fork 3 child streams from the main training stream. Each branch dispatches its 4 ops on its own stream. A join event synchronizes all 3 before the loss kernel. Within a CUDA Graph, this uses `cudaStreamBeginCaptureToGraph` with fork/join topology (CUDA 12+ feature). + +**Applies to:** Forward pass (line 270), f32 forward (line 325), backward pass (line 461) — all 3 branch loops. + +**Validation:** Output logits must be bit-exact vs sequential dispatch. The branches write to non-overlapping regions of `b_logits_buf` (offset by `logit_byte_offset`), so no race conditions. + +**Risk:** CUDA Graph fork/join is complex. If cudarc doesn't expose the necessary APIs, we may need raw CUDA driver calls via `cudarc::driver::sys`. + +### 4. Multi-stream target/online forward parallelism + +**File:** `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs:4307-4323` + +**Current:** Pass 2 (target forward on next_states) and Pass 3 (online forward on next_states for Double DQN) execute sequentially. They read from the same `next_states_buf` but write to completely independent output buffers. + +**Fix:** Dispatch Pass 2 and Pass 3 on separate streams. Join before the loss kernel. Within the CUDA Graph, this is another fork/join. + +**Validation:** Target Q-values and online Q-values must be bit-exact vs sequential. No shared write buffers — safe to parallelize. + +### 5. Dynamic batch sizing via AutoBatchSizer + +**File:** `crates/ml/src/trainers/dqn/trainer/constructor.rs`, `crates/ml-core/src/memory_optimization/auto_batch_size.rs` + +**Current:** `AutoBatchSizer` already computes optimal batch size from VRAM — the H100 log shows `VRAM ceiling = 2,085,808` — but CLI `--batch-size 128` overrides it. The h100.toml sets `batch_size = 1024` which also caps it far below optimal. + +**Fix:** Let AutoBatchSizer drive the batch size. Remove the static h100.toml override. The sizer already accounts for replay buffer VRAM, per-batch tensor allocations, and kernel shared memory. On H100 with 80GB VRAM and 10M buffer (~35GB), the sizer should yield batch sizes of 4096-8192+ for the remaining ~45GB. + +Ensure the C51/MSE loss kernels' shared memory scales correctly with larger batches (shared memory per block is fixed, but grid size grows linearly with batch — OK for H100's 228KB/SM). + +**Validation:** Log the auto-computed batch size. Train 5 epochs at auto size vs 1024, compare loss. Apply linear LR scaling rule if batch increases >4x (LR × batch_ratio). + +### 6. Q-stats async readback + +**File:** `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs:3657` + +**Current:** `reduce_current_q_stats()` calls `memcpy_dtoh` every 50 steps — implicit stream sync. + +**Fix:** Use the same CUDA event pattern. Dispatch the reduction kernel, record an event, read scalars on the next call (50 steps later — always completed by then). + +**Validation:** Q-value stats must match synchronous readback. Logging-only data — no training impact if delayed by 50 steps. + +### 7. Experience collection → training phase overlap + +**File:** `crates/ml/src/trainers/dqn/trainer/training_loop.rs:1196` + +**Current:** `stream.synchronize()` blocks between experience collection and training. The CPU waits for all 204K experience kernels to complete before starting training setup. + +**Fix:** Record an event after the last experience kernel. The CPU proceeds with training setup (lock acquisition, batch pre-sampling, graph preparation). Wait on the event only when the first training step actually needs replay buffer data. + +**Validation:** Replay buffer must be fully populated before the first training step samples from it. The event ensures this without a full stream sync. + +### 8. Remove lock contention in prefetch loop + +**File:** `crates/ml/src/trainers/dqn/trainer/training_loop.rs:1326` + +**Current:** PREFETCH_K=32 loop acquires/releases read/write locks per chunk. Between chunks there's a lock release → re-acquire gap. + +**Fix:** With async events, batch pre-sampling can happen in a dedicated thread that holds the read lock for the entire epoch. The training thread holds the write lock. No contention if the PER buffer is lock-free (GPU-resident segment tree already is). + +**Validation:** No training impact — this is a CPU-side optimization. Verify no deadlocks under stress (run 100 epochs). + +### 9. fxcache PVC alignment + +**File:** `infra/k8s/argo/precompute-features-template.yaml` + +**Current:** Precompute uses `--data-dir /data/futures-baseline`, but `train_baseline_rl` internally resolves to `data_dir/ES.FUT`. Cache key mismatch → falls back to 5min DBN extraction on H100. + +**Fix:** Already committed — precompute template now passes `data-dir/SYMBOL`. Need to regenerate the PVC cache and validate the training binary hits fxcache. + +**Validation:** H100 training log must show "fxcache hit" instead of "Loading OHLCV bars from DBN files". + +--- + +## Implementation Order + +**Phase 1 (biggest impact, lowest risk):** +1. CUDA Events (item 1) — eliminates 341 syncs/epoch +5. Batch size increase (item 5) — config change only +6. Q-stats async (item 6) — small, safe +9. fxcache alignment (item 9) — already committed, just regenerate + +**Phase 2 (medium complexity):** +7. Experience → training overlap (item 7) +8. Lock contention removal (item 8) + +**Phase 3 (high complexity, high reward):** +2. CUDA Graph for experience collection (item 2) +3. Multi-stream branch parallelism (item 3) +4. Multi-stream target/online parallelism (item 4) + +--- + +## Success Criteria + +- Epoch time: <10s on H100 (from 2 min) +- SM utilization: >85% during training (from 52%) +- Memory bandwidth: >50% during GEMM (from 0%) +- TDP: >250W during training (from 102W) +- Loss trajectory: identical to baseline within floating-point tolerance +- No GPU memory leaks over 200 epochs