5.6 KiB
H100 Maximum Performance — Complete GPU Optimization
Problem
The H100 (80GB, 132 SMs, 3TB/s HBM3, SM 9.0) is severely underutilized during DQN training. Systematic audit found 8 categories of waste totaling ~1-2s per epoch (200 epochs = 200-400s wasted per training run).
Findings & Fixes
P0: Blocking (Training hangs)
1. Segment Tree Race Condition
seg_tree_update and seg_tree_insert kernels use non-atomic tree[node] = tree[2*node] + tree[2*node+1] with 8192 threads racing up the tree. On H100's 132 SMs, all threads execute simultaneously → data race → hang or corrupt tree.
Fix: atomicAdd with deltas. Each thread computes delta = new_leaf - old_leaf, then atomicAdd(&tree[node], delta) up to root. Commutative + associative = race-free. O(log n) per thread, hardware-accelerated on SM 9.0.
2. Per-Step event.synchronize() in Readback
replay_adam_and_readback() blocks CPU waiting for previous step's 12-byte async DtoH. On H100, GPU compute per step is long enough that the event is almost always complete, but the synchronize() call has overhead and can stall in burst scenarios.
Fix: Non-blocking readback — skip event.synchronize(), read stale values from pinned buffer if event not complete. The training guard uses GPU-resident buffers directly; CPU readback values are for logging only.
3. Diversity Loss Sync Readback
readback_diversity_loss() does cuStreamSynchronize + sync cuMemcpyDtoH. Once per epoch but full pipeline stall.
Fix: Async DtoH to pinned buffer slot 9, double-buffered like causal readback.
P1: High Impact (~950ms/epoch)
4. Spectral Norm D2D Weight Sync (450ms/epoch)
After spectral normalization, 10 memcpy_dtod_async calls sync weights from per-layer CudaSlice tensors back to the flattened params_buf. Total: ~25MB per step × 1000 steps = 25GB/epoch of pure copy overhead.
Root cause: Weights stored in two places — per-layer slices (for spectral norm) and flattened buffer (for Adam). The copy bridges them.
Fix: Refactor weight storage so per-layer weight references are views (pointer+offset) into params_buf, not separate allocations. Spectral norm operates directly on the flattened buffer at the correct offsets. Zero D2D copies needed.
5. Spectral Norm Launch Storm (500ms/epoch)
6-12 separate single-block kernel launches (grid=1,1,1, block=256) per step for spectral normalization. Each launch wastes 131/132 SMs and incurs ~5-10µs overhead.
Fix: Fuse all spectral norm operations into a single batched kernel. Each block handles one weight matrix's power iteration. Launch with grid=(num_weight_matrices, 1, 1) — all matrices normalized in parallel.
P2: Medium Impact (~200ms/epoch)
6. Single-Block Finalizer Kernels (100ms/epoch)
9 kernels launch with grid_dim=(1,1,1):
grad_norm_finalize— literally 1 thread (block=1,1,1)causal_q_delta_reduce— single block reduction- CQL finalize, pruning mask, branch head norms
Fix: Batch these into composite multi-block kernels. grad_norm_finalize is a single f32→bf16 cast + sqrt — fuse into the grad_norm accumulation kernel itself.
7. Unnecessary memset_zeros (50µs/step = 50ms/epoch)
5 memset_zeros calls per step on buffers fully overwritten by the next kernel:
total_loss_buf(8 bytes)mse_loss_buf(8 bytes)grad_buf(~1MB — the only one that matters)d_value_logits_buf,d_adv_logits_buf
Fix: Initialize accumulators to zero inside the loss/gradient kernels (first-write semantics via thread 0 memset or atomicExch with 0 on first touch). For grad_buf, keep the memset — it's actually needed since backward accumulates via atomicAdd.
8. cuBLAS TF32 for Forward Pass (variable, up to 2× GEMM speedup)
All GEMMs use CUBLAS_COMPUTE_32F (regular F32 accumulation). On H100, CUBLAS_COMPUTE_32F_FAST_TF32 provides 2-3× throughput with ~10-bit mantissa precision.
Fix: Use TF32 for forward-only GEMMs (Q-value computation, target network). Keep F32 for backward GEMMs (gradient computation needs full precision). The forward pass dominates compute time.
P3: Low Impact but Good Practice
9. Dual-Stream Parallelization (~50µs/step)
The DDQN forward graph replays on the main stream despite having a separate double_dqn_stream. No overlap between online forward and DDQN forward.
Fix: Launch graph_forward_ddqn on double_dqn_stream, synchronize via event before graph_adam. The two forward passes run in parallel on separate SMs.
10. CUDA Graph Expansion for Spectral Norm
After fusing spectral norm (fix #5), capture it as a CUDA graph. Reduces per-step overhead from N kernel launches to 1 graph replay.
Files Modified
| File | Changes |
|---|---|
seg_tree_kernel.cu (ml-dqn) |
atomicAdd delta propagation in update + insert |
gpu_dqn_trainer.rs |
Non-blocking readback, diversity async, TF32 forward, memset elimination, single-block fusion, dual-stream DDQN, weight view refactor |
batched_forward.rs |
TF32 compute type for forward GEMMs |
batched_backward.rs |
Keep F32 for backward GEMMs |
dqn_utility_kernels.cu |
Fused spectral norm kernel, fused grad_norm_finalize |
fused_training.rs |
Graph expansion for spectral norm, dual-stream orchestration |
Success Criteria
- Zero
stream.synchronize()in per-step hot loop - Zero
event.synchronize()in per-step hot loop - Zero
memcpy_dtodfor weight sync in per-step hot loop - All spectral norm in 1-2 kernel launches (not 6-12)
- No single-block (grid=1) kernels in hot path
- TF32 enabled for forward GEMMs
- DDQN forward overlapped with online forward
- H100 training completes epoch 1 without hang
- 19/19 smoke tests pass