96 lines
5.6 KiB
Markdown
96 lines
5.6 KiB
Markdown
# 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_dtod` for 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
|