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>
9.5 KiB
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):
- CUDA Events (item 1) — eliminates 341 syncs/epoch
- Batch size increase (item 5) — config change only
- Q-stats async (item 6) — small, safe
- 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