From fdc8f30c8db5ecd3f2f2f2c030dfd98b4cbb869a Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Wed, 8 Apr 2026 08:48:36 +0200 Subject: [PATCH] =?UTF-8?q?docs:=20H100=20graph-level=20optimization=20spe?= =?UTF-8?q?c=20=E2=80=94=204=20deferred=20fixes=20(~13s/epoch)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fix 2: Eliminate validation graph re-capture (DtoD into same buffers) Fix 3: Async validation on separate CUDA stream (overlap with experience) Fix 4: Aux GEMM in graph_aux (GPU scalar + kernel sig change) Fix 7: Mega-graph fusion (spectral+forward+aux → 2 launches) Combined with 5 implemented fixes: 47s → ~16s/epoch Co-Authored-By: Claude Opus 4.6 (1M context) --- .../2026-04-08-h100-graph-optimizations.md | 270 ++++++++++++++++++ 1 file changed, 270 insertions(+) create mode 100644 docs/superpowers/specs/2026-04-08-h100-graph-optimizations.md diff --git a/docs/superpowers/specs/2026-04-08-h100-graph-optimizations.md b/docs/superpowers/specs/2026-04-08-h100-graph-optimizations.md new file mode 100644 index 000000000..02555ef51 --- /dev/null +++ b/docs/superpowers/specs/2026-04-08-h100-graph-optimizations.md @@ -0,0 +1,270 @@ +# H100 Graph-Level Optimizations: 4 Deferred Fixes + +**Goal:** Implement the 4 remaining H100 performance optimizations (Fixes 2,3,4,7) saving ~13s/epoch. Combined with the 5 already-implemented fixes, total epoch time drops from 47s to ~16s. + +**Architecture:** CUDA graph management improvements — eliminate re-capture, async overlap, kernel fusion, mega-graph. All GPU-side, no CPU path. + +**Tech Stack:** CUDA driver API (cuGraphExecUpdate, cuStreamWaitEvent, cuEventRecord), cudarc, cuBLAS + +--- + +## Fix 2: Eliminate Validation Graph Re-capture (saves ~4s/epoch) + +### Problem +`compute_validation_loss()` calls `evaluator.invalidate_dqn_graph()` every epoch (metrics.rs:510), forcing full CUDA graph re-capture of the backtest forward pass. + +### Solution +The evaluator has its OWN weight buffers (`eval_weights`, `eval_branching`) that the graph was captured against. Instead of invalidating, DtoD copy new weights INTO those existing buffers. The graph's captured pointers remain valid — no re-capture needed. + +### Implementation +In `metrics.rs`, replace `evaluator.invalidate_dqn_graph()` with: +```rust +evaluator.update_weights(online_weights, branching_weights)?; +``` + +In `gpu_backtest_evaluator.rs`, add: +```rust +pub fn update_weights( + &mut self, + online: &DuelingWeightSet, + branching: &BranchingWeightSet, +) -> Result<(), MLError> { + // DtoD copy into evaluator's own weight buffers (same addresses graph was captured with) + self.copy_weights_from(online, branching)?; + // No invalidation — topology unchanged, data updated in-place + Ok(()) +} +``` + +The evaluator already has `copy_weights_from` or similar (check for `sync_weights` or `upload_weights` methods). If not, add DtoD copies for each weight tensor using the same pattern as `sync_one` but into the evaluator's buffers. + +### Files +- Modify: `crates/ml/src/trainers/dqn/trainer/metrics.rs` — replace invalidate call +- Modify: `crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs` — add update_weights method + +--- + +## Fix 3: Async Validation on Separate Stream (saves ~8s overlap) + +### Problem +Validation runs synchronously on the main stream, blocking the next epoch's experience collection. + +### Solution +Create a dedicated validation stream + event. Validation runs on stream 2 while the next epoch's experience collection runs on stream 1. + +### Implementation + +**Add fields to DQNTrainer** (`mod.rs`): +```rust +pub(crate) validation_stream: Option>, +pub(crate) validation_event: Option, +pub(crate) pending_val_loss: Option, +``` + +Initialize in constructor: +```rust +validation_stream: if let Some(ref s) = cuda_stream { + Some(Arc::new(s.context().new_stream() + .map_err(|e| anyhow::anyhow!("validation stream: {e}"))?)) +} else { None }, +validation_event: if cuda_stream.is_some() { + let mut event: cudarc::driver::sys::CUevent = std::ptr::null_mut(); + unsafe { cudarc::driver::sys::cuEventCreate(&mut event, 0); } + Some(event) +} else { None }, +pending_val_loss: None, +``` + +**Modify epoch loop** (`training_loop.rs`): +```rust +// At epoch START: collect pending validation from previous epoch +if let Some(val) = self.pending_val_loss.take() { + val_loss = val; +} + +// At epoch END (after training, before next epoch): +// Record event on main stream +if let (Some(ref val_stream), Some(event)) = (&self.validation_stream, self.validation_event) { + let main = self.cuda_stream.as_ref().unwrap(); + unsafe { + cudarc::driver::sys::cuEventRecord(event, main.cu_stream()); + cudarc::driver::sys::cuStreamWaitEvent(val_stream.cu_stream(), event, 0); + } + // Launch validation on stream 2 (non-blocking on main) + let val_result = self.compute_validation_loss().await; + self.pending_val_loss = val_result.ok(); + // Main stream continues to next epoch immediately +} +``` + +**The backtest evaluator** needs to use the validation stream for its cuBLAS calls. Add a `set_stream()` method or pass the stream to `evaluate_dqn_graphed()`. The evaluator's internal cuBLAS handle must be associated with the validation stream. + +If the evaluator creates its own cuBLAS handle at construction, it can be bound to the validation stream at init time. The graph captured on this stream replays on this stream — no cross-stream issues. + +### Files +- Modify: `crates/ml/src/trainers/dqn/trainer/mod.rs` — add fields +- Modify: `crates/ml/src/trainers/dqn/trainer/training_loop.rs` — async launch +- Modify: `crates/ml/src/trainers/dqn/trainer/metrics.rs` — pass validation stream +- Modify: `crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs` — accept stream param + +--- + +## Fix 4: Aux GEMM in graph_aux (saves ~0.5s/epoch) + +### Problem +Exposure aux GEMM runs as 3 separate kernel launches outside graph_aux. The `aux_weight` scalar is passed by value — graph captures the VALUE, not a pointer. + +### Solution +1. Store `aux_weight` in a GPU-resident scalar buffer `CudaSlice` of size 1 +2. Modify `exposure_aux_grad_kernel` signature: `float aux_weight` → `const float* aux_weight_ptr` +3. Before graph_aux replay, async HtoD the current weight value into the scalar buffer +4. Move the aux GEMM call from `run_full_step()` into `submit_aux_ops()` + +### Implementation + +**CUDA kernel change** (`dqn_utility_kernels.cu`): +```cuda +// Change parameter from scalar to pointer: +extern "C" __global__ void exposure_aux_grad_kernel( + ... + const float* __restrict__ aux_weight_ptr, // was: float aux_weight + ... +) { + float aux_weight = *aux_weight_ptr; // device read + ... +} +``` + +**Rust changes** (`gpu_dqn_trainer.rs`): +```rust +// Add field: +exposure_aux_weight_gpu: CudaSlice, // [1] stable address + +// Allocate: +let exposure_aux_weight_gpu = stream.alloc_zeros::(1)?; + +// In launch method: pass pointer instead of value +.arg(&self.exposure_aux_weight_gpu.raw_ptr()) // was: .arg(&aux_weight) +``` + +**Update weight before replay** (`fused_training.rs`): +```rust +// In run_full_step(), BEFORE graph_aux replay: +let weight_val = [self.exposure_aux_weight as f32]; +unsafe { + cudarc::driver::sys::cuMemcpyHtoDAsync_v2( + self.trainer.exposure_aux_weight_gpu.raw_ptr(), + weight_val.as_ptr().cast(), + 4, // sizeof(f32) + self.stream.cu_stream(), + ); +} +``` + +**Move aux GEMM into submit_aux_ops** (`fused_training.rs`): +Cut the aux GEMM block from `run_full_step()` (line ~826-835) and paste at the end of `submit_aux_ops()`. + +### Files +- Modify: `crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu` — kernel sig change +- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` — scalar buffer + launch update +- Modify: `crates/ml/src/trainers/dqn/fused_training.rs` — HtoD update + move aux into submit_aux_ops + +--- + +## Fix 7: Mega-Graph Fusion (saves ~0.3s/epoch) + +### Problem +4 graph launches per training step: spectral → forward → aux → adam = 6ms overhead/step. + +### Solution +Merge spectral + forward + aux into one `graph_mega`. Conditional ops (vaccine, causal) stay outside. Result: 2 launches per step (mega + adam) = 3ms overhead. + +### Implementation + +**Add field** (`fused_training.rs`): +```rust +graph_mega: Option, +``` + +**New capture method:** +```rust +fn capture_graph_mega(&mut self, agent: &mut DQNAgentType, gpu_batch: &GpuBatch) -> Result<()> { + self.trainer.sync_all_streams()?; + let _ = self.stream.context().check_err(); + + let cu_stream = self.stream.cu_stream(); + let mut graph: cuda_sys::CUgraph = std::ptr::null_mut(); + let mut exec: cuda_sys::CUgraphExec = std::ptr::null_mut(); + + unsafe { cuda_sys::cuStreamBeginCapture_v2(cu_stream, THREAD_LOCAL); } + + // Phase 1: Spectral norm + self.trainer.apply_spectral_norm(&self.online_dueling, &self.online_branching)?; + + // Phase 2: Forward + backward (upload + forward + loss + backward) + self.trainer.submit_forward_ops_main()?; + + // Phase 3: Aux ops (clip + EMA + attention + IQL + IQN + CQL + aux GEMM) + self.submit_aux_ops(agent, gpu_batch)?; + + unsafe { cuda_sys::cuStreamEndCapture(cu_stream, &mut graph); } + unsafe { cuda_sys::cuGraphInstantiateWithFlags(&mut exec, graph, 0); } + + self.graph_mega = Some(RawCudaGraph { exec, graph }); + self.graph_spectral = None; // superseded + self.graph_forward = None; // superseded + self.graph_aux = None; // superseded + Ok(()) +} +``` + +**Update run_full_step:** +```rust +if let Some(ref graph) = self.graph_mega { + // Update dynamic values via HtoD (read by graph kernels from device pointers) + self.pre_replay_state_update(agent); + self.update_aux_weight_gpu()?; // Fix 4: aux weight scalar + graph.launch(self.stream.cu_stream())?; +} else { + // Steps 0-1: ungraphed (same as current) + self.trainer.apply_spectral_norm(...)?; + self.trainer.train_step_gpu(...)?; + self.submit_aux_ops(agent, gpu_batch)?; + + // Capture mega-graph at step 2 + if self.steps_since_varmap_sync == 2 { + self.capture_graph_mega(agent, gpu_batch)?; + } +} + +// Steps that stay OUTSIDE the mega-graph: +// - Conditional ops (vaccine, causal) — step-dependent +// - Exposure aux targets DtoD copy — data-dependent +// - compute_grad_norm_outside_graph — already outside +// - replay_adam — separate graph (different frequency possible) +``` + +### Files +- Modify: `crates/ml/src/trainers/dqn/fused_training.rs` — new capture + updated run_full_step + +--- + +## Combined Impact + +| Fix | Savings | Mechanism | +|-----|---------|-----------| +| 2: No re-capture | -4s/epoch | DtoD into existing buffers | +| 3: Async validation | -8s/epoch | Separate stream overlap | +| 4: Aux in graph | -0.5s/epoch | GPU scalar + kernel sig change | +| 7: Mega-graph | -0.3s/epoch | 4 launches → 2 launches | +| **Total** | **-12.8s/epoch** | | + +With the 5 already-implemented fixes (-18s), total savings: **-30.8s/epoch** (47s → ~16s). + +50 trials × 40 epochs × 16s = **8.9 hours** (was 26 hours — 66% reduction, $51 saved per run) + +## Testing +- `cargo test -p ml --lib` — 900+ pass, 0 regressions +- GPU smoke test on RTX 3050 — no SIGSEGV from graph changes +- Verify validation still produces correct Sharpe (async doesn't corrupt) +- Verify mega-graph replay matches ungraphed execution (step 0 vs step 3+)