0b185ddb54
CUDA events, multi-stream branches, graph-captured experience collection, dynamic batch sizing, async readbacks, fxcache alignment. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
152 lines
9.5 KiB
Markdown
152 lines
9.5 KiB
Markdown
# 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
|