From 171fb0832cb82566e99713698f57bb2932b3ad1a Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Wed, 1 Apr 2026 21:41:21 +0200 Subject: [PATCH] =?UTF-8?q?docs:=20H100=20GPU=20optimization=20implementat?= =?UTF-8?q?ion=20plan=20=E2=80=94=2010=20tasks?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Phase 1: CUDA events, dynamic batch size, async q-stats, phase overlap Phase 2: Multi-stream branches, target/online parallelism Phase 3: CUDA Graph experience collection, lock removal Validation: H100 benchmark with epoch time + SM utilization targets Co-Authored-By: Claude Opus 4.6 (1M context) --- .../plans/2026-04-01-h100-gpu-optimization.md | 621 ++++++++++++++++++ 1 file changed, 621 insertions(+) create mode 100644 docs/superpowers/plans/2026-04-01-h100-gpu-optimization.md diff --git a/docs/superpowers/plans/2026-04-01-h100-gpu-optimization.md b/docs/superpowers/plans/2026-04-01-h100-gpu-optimization.md new file mode 100644 index 000000000..086cdde60 --- /dev/null +++ b/docs/superpowers/plans/2026-04-01-h100-gpu-optimization.md @@ -0,0 +1,621 @@ +# H100 GPU Pipeline Optimization Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Reduce DQN epoch time from 2 min to <10s on H100 by eliminating GPU idle time from blocking synchronize calls, maximizing batch size, and enabling multi-stream parallelism. + +**Architecture:** Replace `stream.synchronize()` with CUDA events for async scalar readback. Let `AutoBatchSizer` drive batch size (removes static cap). Capture experience collection as a CUDA Graph. Fork child streams for independent branch/target forward passes. + +**Tech Stack:** Rust, cudarc 0.17 (CudaEvent, CudaStream::fork), CUDA 12.4, H100 PCIe 80GB. + +--- + +## File Structure + +| File | Action | Change | +|------|--------|--------| +| `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` | Modify | CUDA events in `replay_adam_and_readback`, async q-stats, multi-stream forward | +| `crates/ml/src/cuda_pipeline/batched_forward.rs` | Modify | Multi-stream branch dispatch | +| `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` | Modify | CUDA Graph capture for timestep loop | +| `crates/ml/src/trainers/dqn/trainer/training_loop.rs` | Modify | Async training loop, event-based phase transitions | +| `config/gpu/h100.toml` | Modify | Remove static batch_size cap | +| `crates/ml/src/trainers/dqn/trainer/constructor.rs` | Modify | Let AutoBatchSizer drive batch_size | +| `infra/k8s/argo/precompute-features-template.yaml` | Done | Already committed (data-dir/SYMBOL alignment) | + +--- + +## Phase 1: Eliminate GPU Idle Time (biggest impact) + +### Task 1: CUDA Events replace per-step stream.synchronize() + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs:2865-2893` + +The `replay_adam_and_readback()` method calls `stream.synchronize()` on every training step (341×/epoch). This is the primary bottleneck — H100 finishes both CUDA graphs in microseconds then sits idle. + +- [ ] **Step 1: Add event and async readback state to GpuDqnTrainer struct** + +In `gpu_dqn_trainer.rs`, add fields to the `GpuDqnTrainer` struct (near line 479): + +```rust +/// CUDA event for async scalar readback (replaces stream.synchronize per step) +readback_event: Option>, +/// Host-side buffers for deferred readback (loss, grad_norm from previous step) +pending_loss: f32, +pending_grad_norm: f32, +/// Whether there's an in-flight readback to collect +readback_pending: bool, +``` + +Initialize in `new()`: +```rust +readback_event: None, +pending_loss: 0.0, +pending_grad_norm: 0.0, +readback_pending: false, +``` + +- [ ] **Step 2: Rewrite replay_adam_and_readback to use events** + +Replace the synchronous `replay_adam_and_readback` method: + +```rust +/// Collect scalars from the PREVIOUS step (non-blocking poll), +/// then launch grad_norm + adam for the CURRENT step and record an event. +pub fn replay_adam_and_readback(&mut self) -> Result { + // 1. Collect previous step's scalars (if any) + let prev_scalars = if self.readback_pending { + // Poll: is the previous step's event done? + if let Some(ref event) = self.readback_event { + if !event.is_complete() { + // Still in-flight — wait on just this event (not full stream) + event.synchronize().map_err(|e| + MLError::ModelError(format!("readback event sync: {e}")) + )?; + } + } + // Now safe to read the scalars from the previous step + self.readback_pending = false; + FusedTrainScalars { + total_loss: self.pending_loss, + grad_norm: self.pending_grad_norm, + } + } else { + FusedTrainScalars { total_loss: 0.0, grad_norm: 0.0 } + }; + + // 2. Launch current step: grad_norm + adam + self.compute_grad_norm_outside_graph()?; + self.replay_adam()?; + + // 3. Issue async DtoH for current step's scalars + unsafe { + cudarc::driver::sys::cuMemcpyDtoHAsync_v2( + (&mut self.pending_loss as *mut f32).cast(), + self.total_loss_buf.raw_ptr(), + std::mem::size_of::(), + *self.stream.cu_stream(), + ); + cudarc::driver::sys::cuMemcpyDtoHAsync_v2( + (&mut self.pending_grad_norm as *mut f32).cast(), + self.grad_norm_f32_buf.raw_ptr(), + std::mem::size_of::(), + *self.stream.cu_stream(), + ); + } + + // 4. Record event marking when the async DtoH completes + self.readback_event = Some(self.stream.record_event(None).map_err(|e| + MLError::ModelError(format!("readback event record: {e}")) + )?); + self.readback_pending = true; + + // Return PREVIOUS step's scalars (current step's will be collected next call) + self.scalars_readback_host = [prev_scalars.total_loss, prev_scalars.grad_norm]; + Ok(prev_scalars) +} + +/// Flush: block until the last in-flight readback completes. +/// Called once at epoch end. +pub fn flush_readback(&mut self) -> Result { + if !self.readback_pending { + return Ok(FusedTrainScalars { total_loss: 0.0, grad_norm: 0.0 }); + } + if let Some(ref event) = self.readback_event { + if !event.is_complete() { + event.synchronize().map_err(|e| + MLError::ModelError(format!("flush readback sync: {e}")) + )?; + } + } + self.readback_pending = false; + Ok(FusedTrainScalars { + total_loss: self.pending_loss, + grad_norm: self.pending_grad_norm, + }) +} +``` + +- [ ] **Step 3: Update training loop to flush at epoch end** + +In `training_loop.rs`, after the training step loop (after the `for chunk_start in ...` loop ends), add: + +```rust +// Flush the last in-flight readback +if let Some(ref mut fused) = self.fused_ctx { + let final_scalars = fused.flush_readback()?; + // Use final_scalars for epoch-end logging if needed +} +``` + +- [ ] **Step 4: Verify compile** + +Run: `SQLX_OFFLINE=true cargo check -p ml --lib` +Expected: compiles clean + +- [ ] **Step 5: Run local smoke test to validate loss trajectory** + +Run: `FOXHUNT_TEST_DATA=test_data/futures-baseline FOXHUNT_FEATURE_CACHE_DIR=test_data/feature-cache SQLX_OFFLINE=true cargo test -p ml --lib -- smoke_tests::training_stability::test_50_epoch_convergence --ignored --nocapture` + +Expected: test passes, loss values are finite, Sharpe is reasonable. Loss trajectory may differ slightly from synchronous path (1-step lag on PER priorities) but should converge similarly. + +- [ ] **Step 6: Commit** + +```bash +git commit -m "perf: CUDA events replace per-step stream.synchronize in training loop" +``` + +--- + +### Task 2: Async Q-stats readback + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs:3640-3660` + +The `reduce_current_q_stats()` calls `memcpy_dtoh` every 50 steps — implicit sync. Use the same event pattern. + +- [ ] **Step 1: Add q-stats event and pending buffer** + +Add to `GpuDqnTrainer` struct: + +```rust +/// Pending q-stats from previous reduction (deferred readback) +q_stats_event: Option>, +q_stats_pending: [f32; 5], +q_stats_ready: bool, +``` + +Initialize: +```rust +q_stats_event: None, +q_stats_pending: [0.0; 5], +q_stats_ready: false, +``` + +- [ ] **Step 2: Rewrite reduce_current_q_stats for async readback** + +Replace the synchronous `memcpy_dtoh` in `reduce_current_q_stats()`: + +```rust +pub fn reduce_current_q_stats(&mut self) -> Result { + // Collect previous reduction if ready + let result = if self.q_stats_ready { + if let Some(ref event) = self.q_stats_event { + if !event.is_complete() { + event.synchronize().map_err(|e| + MLError::ModelError(format!("q_stats event sync: {e}")) + )?; + } + } + self.q_stats_ready = false; + let h = self.q_stats_pending; + QValueStatsResult { + avg_max_q: h[0] as f64, q_min: h[1], q_max: h[2], q_mean: h[3], q_var: h[4], + } + } else { + QValueStatsResult { avg_max_q: 0.0, q_min: 0.0, q_max: 0.0, q_mean: 0.0, q_var: 0.0 } + }; + + // Launch new reduction kernel + // ... (existing kernel launch code stays the same) ... + + // Async readback + unsafe { + cudarc::driver::sys::cuMemcpyDtoHAsync_v2( + self.q_stats_pending.as_mut_ptr().cast(), + self.q_stats_buf.raw_ptr(), + 5 * std::mem::size_of::(), + *self.stream.cu_stream(), + ); + } + self.q_stats_event = Some(self.stream.record_event(None).map_err(|e| + MLError::ModelError(format!("q_stats event record: {e}")) + )?); + self.q_stats_ready = true; + + Ok(result) +} +``` + +- [ ] **Step 3: Verify compile + smoke test** + +Run: `SQLX_OFFLINE=true cargo check -p ml --lib` +Run: smoke test (same as Task 1 Step 5) + +- [ ] **Step 4: Commit** + +```bash +git commit -m "perf: async q-stats readback via CUDA events" +``` + +--- + +### Task 3: Dynamic batch sizing — remove static cap + +**Files:** +- Modify: `config/gpu/h100.toml:3` +- Modify: `crates/ml/src/trainers/dqn/trainer/constructor.rs` (AutoBatchSizer logic) + +- [ ] **Step 1: Remove static batch_size from h100.toml** + +Change `config/gpu/h100.toml` line 3: + +```toml +# batch_size is auto-computed by AutoBatchSizer from available VRAM. +# Do NOT set a static value here — let the sizer maximize GPU utilization. +# batch_size = 1024 # REMOVED: was capping H100 far below optimal +``` + +Actually: the TOML profile must still have the field for the parser. Set it to 0 as a sentinel meaning "auto": + +```toml +batch_size = 0 # 0 = auto-compute from VRAM (AutoBatchSizer) +``` + +- [ ] **Step 2: Update constructor to honor batch_size=0 as auto** + +In `constructor.rs`, find where `hyperparams.batch_size` is set from the profile. After the profile applies, if `batch_size == 0`, let AutoBatchSizer compute it: + +```rust +// After profile.apply_to(&mut hyperparams): +if hyperparams.batch_size == 0 { + // Auto-compute: AutoBatchSizer already ran and set the VRAM ceiling. + // Use the ceiling as the batch size, capped at a sane max (16384). + hyperparams.batch_size = auto_batch_ceiling.min(16384); + info!("AutoBatchSizer: batch_size auto-computed to {}", hyperparams.batch_size); +} +``` + +- [ ] **Step 3: Verify the auto-computed value is reasonable** + +Run locally (RTX 3050): the sizer should still pick a small batch (64-128) for 4GB VRAM. +Check H100 would get: look at the `AutoBatchSizer: GPU VRAM ceiling = 2085808` log — that's the theoretical max. The actual batch should be much smaller (constrained by kernel shared memory and practical gradient noise). Cap at 8192 or 16384. + +- [ ] **Step 4: Commit** + +```bash +git commit -m "perf: dynamic batch sizing — AutoBatchSizer drives batch_size on H100" +``` + +--- + +### Task 4: Experience collection → training phase overlap + +**Files:** +- Modify: `crates/ml/src/trainers/dqn/trainer/training_loop.rs:1196` + +- [ ] **Step 1: Replace stream.synchronize() with event after experience collection** + +At line 1196, change: + +```rust +// BEFORE: +stream.synchronize() + .map_err(|e| MLError::DeviceError(format!("experience sync: {e}")))?; + +// AFTER: +let experience_done_event = stream.record_event(None) + .map_err(|e| MLError::DeviceError(format!("experience event record: {e}")))?; +``` + +Then before the first training step that samples from the replay buffer, wait on the event: + +```rust +// Before first replay buffer sample: +if !experience_done_event.is_complete() { + experience_done_event.synchronize() + .map_err(|e| MLError::DeviceError(format!("experience event wait: {e}")))?; +} +``` + +This allows CPU-side training setup (lock acquisition, pre-sampling initialization) to overlap with the tail end of experience collection. + +- [ ] **Step 2: Verify + commit** + +```bash +git commit -m "perf: event-based experience→training phase transition" +``` + +--- + +### Task 5: fxcache PVC regeneration + +- [ ] **Step 1: Regenerate PVC cache with aligned data-dir** + +The precompute template already uses `data-dir/SYMBOL`. Just re-run: + +```bash +./scripts/argo-precompute.sh ES.FUT --pool ci-compile-cpu +``` + +- [ ] **Step 2: Verify training hits fxcache** + +Submit a quick 2-epoch baseline and check logs for "fxcache hit": + +```bash +argo submit -n foxhunt --from=wftmpl/train-baseline-rl -p train-epochs=2 +``` + +Check: `argo logs | grep fxcache` + +- [ ] **Step 3: Commit (if any template changes needed)** + +--- + +## Phase 2: Multi-Stream Parallelism + +### Task 6: Multi-stream branch dispatch in forward pass + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/batched_forward.rs:265-284` +- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` (add child streams) + +- [ ] **Step 1: Add 3 child streams to CublasForward** + +In `batched_forward.rs`, add to the struct: + +```rust +/// Child streams for parallel branch dispatch (forked from main training stream) +branch_streams: [Arc; 3], +``` + +Create them in the constructor by forking: + +```rust +let branch_streams = [ + stream.fork().map_err(|e| MLError::DeviceError(format!("fork branch stream 0: {e}")))?, + stream.fork().map_err(|e| MLError::DeviceError(format!("fork branch stream 1: {e}")))?, + stream.fork().map_err(|e| MLError::DeviceError(format!("fork branch stream 2: {e}")))?, +]; +``` + +- [ ] **Step 2: Dispatch branches on separate streams** + +In `forward_online_raw()`, replace the sequential `for d in 0..3` loop (lines 270-284): + +```rust +// Fork: dispatch each branch on its own stream +for d in 0..3 { + let branch_stream = &self.branch_streams[d]; + // Wait for trunk (h_s2) to be ready on this branch stream + let trunk_event = self.stream.record_event(None) + .map_err(|e| MLError::ModelError(format!("trunk event: {e}")))?; + branch_stream.wait(&trunk_event) + .map_err(|e| MLError::ModelError(format!("branch wait trunk: {e}")))?; + + let n_d = branch_sizes[d]; + let w_fc_idx = branch_w_base[d]; + + // Hidden layer + self.gemmex_bf16_on_stream(branch_stream, w_ptrs[w_fc_idx], h_s2_ptr, branch_h_ptrs[d], self.adv_h, b, self.shared_h2, "h_bd")?; + self.launch_add_bias_relu_bf16_on_stream(branch_stream, branch_h_ptrs[d], w_ptrs[w_fc_idx + 1], self.adv_h, b)?; + + // Output layer + let adv_out_ptr = b_logits_ptr + logit_byte_offset; + self.gemmex_bf16_to_f32_on_stream(branch_stream, w_ptrs[w_fc_idx + 2], branch_h_ptrs[d], adv_out_ptr, n_d * na, b, self.adv_h, "adv_logits")?; + self.launch_add_bias_f32_on_stream(branch_stream, adv_out_ptr, w_ptrs[w_fc_idx + 3], n_d * na, b)?; + + logit_byte_offset += (b * n_d * na * std::mem::size_of::()) as u64; +} + +// Join: main stream waits for all branches to finish +for d in 0..3 { + let branch_done = self.branch_streams[d].record_event(None) + .map_err(|e| MLError::ModelError(format!("branch done event: {e}")))?; + self.stream.wait(&branch_done) + .map_err(|e| MLError::ModelError(format!("join branch {d}: {e}")))?; +} +``` + +Note: `gemmex_bf16_on_stream` and `launch_add_bias_relu_bf16_on_stream` are new variants that take an explicit stream parameter. Implement them by extracting the cuBLAS handle → set stream → call GemmEx → restore stream. cuBLAS handles are NOT stream-specific — `cublasSetStream()` switches the active stream per-call. + +- [ ] **Step 3: Apply same pattern to f32 forward (line 325) and backward (line 461)** + +Repeat the fork/join pattern for both the f32 experience collection forward and the backward pass branch loops. + +- [ ] **Step 4: Verify compile + smoke test** + +Run smoke tests — output logits must be identical (branches write to non-overlapping memory regions). + +- [ ] **Step 5: Commit** + +```bash +git commit -m "perf: multi-stream branch dispatch — 3 advantage heads in parallel" +``` + +--- + +### Task 7: Multi-stream target/online forward parallelism + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs:4307-4323` + +- [ ] **Step 1: Add a second stream for Double DQN pass** + +Add to `GpuDqnTrainer`: + +```rust +/// Second stream for parallel target/online next-state forward passes +double_dqn_stream: Arc, +``` + +Fork in constructor. + +- [ ] **Step 2: Dispatch Pass 2 and Pass 3 on separate streams** + +In `submit_forward_ops()`, after Pass 1 (online forward on current states): + +```rust +// Pass 2 + Pass 3 are independent: both read next_states, write to different outputs. +// Fork: record event after stochastic depth (Pass 1 done) +let pass1_done = self.stream.record_event(None)?; +self.double_dqn_stream.wait(&pass1_done)?; + +// Pass 2: target forward on next_states (main stream) +cublas.forward_target_raw( + &self.stream, self.ptrs.next_states_buf, &tg_w_ptrs, ... +)?; + +// Pass 3: online forward on next_states (double_dqn_stream) +cublas.forward_online_raw( + &self.double_dqn_stream, self.ptrs.next_states_buf, &on_w_ptrs, ... +)?; + +// Join before loss kernel +let target_done = self.stream.record_event(None)?; +let online_done = self.double_dqn_stream.record_event(None)?; +self.stream.wait(&online_done)?; +// main stream now has both pass 2 + pass 3 complete +``` + +- [ ] **Step 3: Handle CUDA Graph capture with multi-stream** + +CUDA Graph capture with fork/join requires `CU_STREAM_CAPTURE_MODE_THREAD_LOCAL` (already used at line 3734). The fork/join events within the graph are captured as graph dependencies. Verify this works with cudarc — may need raw `cuStreamWaitEvent` calls during capture. + +- [ ] **Step 4: Verify + smoke test + commit** + +```bash +git commit -m "perf: parallel target/online forward on separate CUDA streams" +``` + +--- + +### Task 8: CUDA Graph for experience collection + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs:1506` + +This is the most complex task. The 100-iteration host loop dispatches 400+ kernels. Capturing as a graph eliminates all host→GPU latency. + +- [ ] **Step 1: Make per-timestep parameters GPU-resident** + +The loop currently passes `epsilon`, step counter, and state pointers from the host. These must become GPU buffers: + +```rust +/// GPU-resident parameters for graph-captured timestep loop +epsilon_buf: CudaSlice, // [1] — updated via DtoH before graph launch +step_counter_buf: CudaSlice, // [1] — incremented by a tiny CUDA kernel per step +``` + +- [ ] **Step 2: Create a step-increment kernel** + +Tiny kernel that increments `step_counter_buf[0]` and computes `epsilon = eps_end + (eps_start - eps_end) * exp(-step / decay)` on GPU. No host involvement. + +- [ ] **Step 3: Capture the 100-timestep loop as a single graph** + +```rust +stream.begin_capture(CU_STREAM_CAPTURE_MODE_THREAD_LOCAL)?; +for t in 0..timesteps { + self.launch_experience_state_gather()?; + self.launch_cublas_forward_online()?; + self.launch_compute_expected_q()?; + self.launch_experience_action_select()?; // reads epsilon_buf + self.launch_experience_env_step()?; + self.launch_step_increment()?; // increments step_counter, updates epsilon +} +let graph = stream.end_capture(CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH)?; +``` + +Before each epoch: `memcpy_htod` the initial epsilon and step counter, then `graph.launch()`. + +- [ ] **Step 4: Handle dynamic epsilon** + +Epsilon decays across episodes within an epoch. The step-increment kernel computes the new epsilon from the step counter using the exponential decay formula, writing to `epsilon_buf`. The action selection kernel reads `epsilon_buf` directly — no host involvement. + +- [ ] **Step 5: Verify collected experiences match sequential path** + +Run 1 epoch with both paths (graph and sequential), compare states/actions/rewards for the first 100 experiences. Must be numerically identical given the same Philox RNG seed. + +- [ ] **Step 6: Commit** + +```bash +git commit -m "perf: CUDA Graph capture for 100-timestep experience collection loop" +``` + +--- + +### Task 9: Lock contention removal + +**Files:** +- Modify: `crates/ml/src/trainers/dqn/trainer/training_loop.rs:1336-1360` + +- [ ] **Step 1: Pre-sample all batches for the epoch in one lock acquisition** + +Instead of acquiring/releasing the read lock every 32 steps, sample ALL batches at epoch start: + +```rust +let all_batches = { + let agent = self.agent.read().await; + let buffer = agent.memory(); + (0..num_training_steps) + .map(|_| buffer.sample(self.current_batch_size)) + .collect::>() +}; +// Lock released. Training loop iterates over pre-sampled batches. +``` + +This is safe because the GPU PER buffer is immutable during training (new experiences are only added in Phase 2, which is finished before training starts). + +- [ ] **Step 2: Remove PREFETCH_K chunking** + +The `for chunk_start in (0..num_training_steps).step_by(PREFETCH_K)` loop becomes a simple `for batch in all_batches.iter()`. + +- [ ] **Step 3: Verify + commit** + +```bash +git commit -m "perf: pre-sample all training batches in single lock acquisition" +``` + +--- + +## Validation + +### Task 10: H100 benchmark + +- [ ] **Step 1: Rebuild and upload binary to PVC** + +Compile in-cluster via the baseline template. + +- [ ] **Step 2: Run 10-epoch baseline and measure** + +```bash +argo submit -n foxhunt --from=wftmpl/train-baseline-rl -p train-epochs=10 +``` + +Check: +- Epoch time: target <10s (from 2 min) +- `nvidia-smi`: SM utilization >85%, memory bandwidth >50%, TDP >250W +- Loss trajectory: finite, converging +- "fxcache hit" in logs + +- [ ] **Step 3: Compare loss with pre-optimization baseline** + +The 7-epoch metrics from the terminated run serve as reference: +- Epoch 4: Sharpe 0.06, PF 0.94 +- Epoch 6: Sharpe -1.12, PF 1.17 + +Post-optimization should show similar or better convergence (larger batches may improve gradient quality). + +- [ ] **Step 4: Run full 50-epoch baseline if metrics look good** + +- [ ] **Step 5: Save metrics to memory**