docs: H100 epoch optimization spec (37s → <5s) + phase profiling
Spec: 3-phase plan to reduce DQN training epoch from 37s to <5s on H100. - Phase 1: eliminate 300 CPU roundtrips in PER sampling (GPU Philox RNG) - Phase 2: increase kernel occupancy (batch_size 512+, cuBLAS profiling) - Phase 3: pipeline overlap (dual-stream experience/training) Profiling instrumentation added to training loop (init/experience/training/validation breakdown). RTX 3050 baseline: training=97.2%, experience=2.3% — PER sampling with CPU RNG + DtoH sync is the primary bottleneck. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -0,0 +1,196 @@
|
||||
# H100 Epoch Optimization: 37s → <5s per epoch
|
||||
|
||||
**Date**: 2026-03-21
|
||||
**Target**: <5s/epoch on H100 (80GB), <10s acceptable, <3s stretch goal
|
||||
**Current**: 37.56s/epoch (H100, 51 atoms, 300 training steps, CUDA Graph)
|
||||
|
||||
## Profiling Results
|
||||
|
||||
RTX 3050 phase breakdown (proportions apply to H100):
|
||||
|
||||
| Phase | Time | % | Per-step |
|
||||
|-------|-----:|--:|---------|
|
||||
| Init (data upload) | 10ms | 0.1% | once/epoch |
|
||||
| Experience collection | 348ms | 2.3% | 1 kernel launch |
|
||||
| **Training steps** | **14,858ms** | **97.2%** | 300 steps |
|
||||
| Validation | 27ms | 0.2% | 1 backtest pass |
|
||||
|
||||
H100 epoch time: 37.56s × 97.2% training = **36.5s in training steps**.
|
||||
At 300 steps/epoch: **122ms per training step**.
|
||||
|
||||
## Root Cause Analysis
|
||||
|
||||
### Primary bottleneck: PER sampling CPU roundtrips (300×/epoch)
|
||||
|
||||
Each `sample_proportional()` call does:
|
||||
1. **DtoH**: `cs_total()` reads prefix sum total (1 scalar) — forces stream sync
|
||||
2. **CPU RNG**: `rand::thread_rng().gen::<f32>()` × batch_size — serial CPU work
|
||||
3. **HtoD**: uploads random thresholds (batch_size × 4 bytes)
|
||||
4. **HtoD**: zeros `batch_max` (4 bytes)
|
||||
|
||||
300 calls × ~100μs per sync barrier = **30ms just for sync**, but the CPU RNG + alloc + memcpy overhead adds 100-400ms per call.
|
||||
|
||||
### Secondary: kernel underutilization
|
||||
|
||||
Training kernel: `block_dim=(32,1,1)`, `grid_dim=(batch_size,1,1)` = 256 blocks × 32 threads = **8,192 threads on 132-SM H100** (3.0% occupancy). Each SM can run 2048 threads — we're using 1 warp per SM on a fraction of 132 SMs.
|
||||
|
||||
### Tertiary: sequential phase execution
|
||||
|
||||
Experience collection → Training → Validation runs serially. No overlap.
|
||||
|
||||
## Design: Three-Phase Optimization
|
||||
|
||||
### Phase 1: Eliminate CPU Roundtrips (target: 10-15s/epoch)
|
||||
|
||||
#### 1a. GPU-native PER threshold generation
|
||||
|
||||
Replace CPU `rand::thread_rng()` with Philox CUDA RNG kernel (already have this pattern in IQN tau sampler):
|
||||
|
||||
```cuda
|
||||
// New kernel: per_generate_thresholds_kernel
|
||||
__global__ void per_generate_thresholds(
|
||||
float* __restrict__ thresholds, // [batch_size] output
|
||||
float total_sum, // from prefix sum (stays on GPU)
|
||||
unsigned long long seed, // counter-based, no state
|
||||
int batch_size)
|
||||
{
|
||||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (i >= batch_size) return;
|
||||
// Philox-based PRNG (same as iqn_sample_taus_kernel)
|
||||
unsigned int r = philox_uint(seed, i);
|
||||
float u = __uint_as_float(0x3f800000u | (r >> 9)) - 1.0f; // uniform [0,1)
|
||||
thresholds[i] = u * total_sum;
|
||||
}
|
||||
```
|
||||
|
||||
**Key change**: `total_sum` stays on GPU — read it via a single DtoD copy from `cs_buf[n-1]` to a scalar buffer, then pass as kernel arg. Zero DtoH.
|
||||
|
||||
#### 1b. Fuse PER kernels into single launch
|
||||
|
||||
Current: 10+ kernel launches per sample (pow_alpha, prefix_sum, searchsorted, 5 gathers, IS weights, reduce).
|
||||
|
||||
Fuse into 2 kernels:
|
||||
1. `per_sample_and_gather_kernel`: prefix sum total (already computed) → thresholds → searchsorted → gather all fields
|
||||
2. `per_is_weights_kernel`: compute IS weights + normalize
|
||||
|
||||
Or capture PER sampling into a second CUDA Graph.
|
||||
|
||||
#### 1c. Remove adam_step HtoD
|
||||
|
||||
Store adam_step as a GPU scalar, increment via 1-element kernel or atomicAdd.
|
||||
|
||||
#### 1d. Audit and remove all hot-path `to_host()` calls
|
||||
|
||||
Remove `GpuTensor::to_host()` from the public API. Replace with `to_host_checkpoint()` (clearly marked for cold-path use only). Compile-time guard via `#[deprecated]` or feature gate.
|
||||
|
||||
### Phase 2: Increase Kernel Occupancy (target: 5-8s/epoch)
|
||||
|
||||
#### 2a. Batched training: multiple samples per warp
|
||||
|
||||
Current: 1 warp (32 threads) per sample, 64 samples per graph replay.
|
||||
Target: 4-8 samples per block (128-256 threads), or increase batch_size to 512-1024.
|
||||
|
||||
With batch_size=512: `grid=(512,1,1), block=(32,1,1)` = 16K threads = 7.5% occupancy.
|
||||
With batch_size=1024 + 4 samples/block: 8K threads at 128 threads/block = better SM utilization.
|
||||
|
||||
#### 2b. cuBLAS for network layers
|
||||
|
||||
Replace warp-cooperative matvec with cuBLAS `cublasSgemmStridedBatched` for the shared trunk layers (256×256 matmul). For batch_size=1024, this is a `[1024, 256] × [256, 256]` GEMM — cuBLAS tensor cores will dominate.
|
||||
|
||||
Trade-off: cuBLAS adds launch overhead but uses tensor cores (3.9 PFLOPS BF16 on H100 vs ~100 TFLOPS custom kernel).
|
||||
|
||||
Profile both and pick winner. Custom kernel may win for batch<256.
|
||||
|
||||
#### 2c. BF16 throughout (not just storage)
|
||||
|
||||
Current: BF16 weight storage → F32 accumulation in kernel.
|
||||
Target: BF16 weights + BF16 activations + F32 accumulation (tensor core matmul).
|
||||
|
||||
H100 tensor cores: BF16 GEMM at 989 TFLOPS vs F32 at 67 TFLOPS = **14.8× potential**.
|
||||
|
||||
### Phase 3: Pipeline Overlap (target: <5s/epoch)
|
||||
|
||||
#### 3a. Dual-stream: overlap experience + training
|
||||
|
||||
```
|
||||
Stream A: [experience_epoch_N+1] → [experience_epoch_N+2] → ...
|
||||
Stream B: [train_epoch_N] → [train_epoch_N+1] → ...
|
||||
↑ uses experiences from epoch N (already in buffer)
|
||||
```
|
||||
|
||||
After epoch 1, experience collection for epoch N+1 overlaps with training for epoch N. ~2× throughput.
|
||||
|
||||
#### 3b. Fuse PER sampling + training into mega-graph
|
||||
|
||||
Capture the entire sequence (PER sample → batch upload → forward+loss → backward → adam → EMA) into a single CUDA Graph. One `cuGraphLaunch` per step instead of 15+ kernel launches.
|
||||
|
||||
Requires: GPU-native PER (Phase 1a) and graph-compatible memory pool.
|
||||
|
||||
#### 3c. Async validation
|
||||
|
||||
Move validation backtest to a separate stream, overlapping with the start of the next epoch's experience collection.
|
||||
|
||||
## Implementation Order
|
||||
|
||||
1. **Phase 1a**: GPU Philox RNG for PER thresholds (eliminate sync barriers)
|
||||
2. **Phase 2a**: Increase batch_size to 512+ on H100 (config change, immediate win, amplifies all later optimizations)
|
||||
3. **Phase 1b**: Fuse PER kernels + pre-allocate all PER buffers (zero `cuMemAlloc` after warmup)
|
||||
4. **Phase 1c**: GPU adam_step increment
|
||||
5. **Phase 1d**: Remove `to_host()` from hot paths, deprecate the function
|
||||
6. **Phase 2b**: Profile cuBLAS vs custom for training matmul
|
||||
7. **Phase 2c**: BF16 activations (with C51 log_softmax numerical guard)
|
||||
8. **Phase 3a**: Dual-stream overlap (requires resolving stream unification tension)
|
||||
9. **Phase 3b**: Mega-graph (if 3a insufficient)
|
||||
|
||||
## Expected Speedup Stack
|
||||
|
||||
| Optimization | Expected speedup | Cumulative |
|
||||
|-------------|:----------------:|:----------:|
|
||||
| GPU PER RNG (eliminate 300 sync barriers) | 2-3× | 12-18s |
|
||||
| Batch_size 512+ (occupancy) | 1.3-1.5× | 9-12s |
|
||||
| Fuse PER kernels + pre-alloc (reduce 14→2 launches) | 1.2-1.5× | 6-10s |
|
||||
| cuBLAS/tensor core matmul | 1.5-3× | 3-6s |
|
||||
| Dual-stream overlap | 1.5-2× | 2-4s |
|
||||
|
||||
Conservative estimate: **5-8s/epoch**. Aggressive: **3-4s/epoch**.
|
||||
|
||||
## Known Risks
|
||||
|
||||
### CUDA Graph capture requires fixed buffer addresses
|
||||
PER's `sample_proportional` currently allocates new `CudaSlice` buffers per call (lines 434, 436, 443, 450-451, 478, 490 — ~10 allocs × 300 calls = 3,000 GPU malloc/free per epoch). For mega-graph capture (Phase 3b), ALL buffers must be pre-allocated at fixed addresses. Phase 1b must pre-allocate these.
|
||||
|
||||
### Dual-stream conflicts with stream unification
|
||||
The `stream_unification_plan.md` explicitly advocates collapsing to a single forked stream to avoid cudarc event tracking conflicts. Phase 3a proposes re-introducing dual streams. Resolution: use separate `CudaContext` per stream (each context has independent event tracking), or complete cudarc event tracking fix first.
|
||||
|
||||
### C51 log_softmax underflow in BF16
|
||||
Phase 2c (BF16 activations) risks numerical issues in C51 distributional loss. `log_softmax` on BF16 can underflow for low-probability atoms. Mitigation: keep loss computation in F32 accumulation, only use BF16 for matmul operands. Profile numerically before deploying.
|
||||
|
||||
### Prefix sum single-block scalability
|
||||
The prefix sum kernel runs with `grid_dim=(1,1,1)` — single block processing up to ~1M elements. At large replay buffer capacities (1M+ on H100), this serializes. May need multi-block parallel scan (CUB `DeviceScan`) for Phase 2a batch_size increases.
|
||||
|
||||
## Success Criteria
|
||||
|
||||
- [ ] Phase breakdown log shows zero DtoH in training steps (except 8-byte loss+grad_norm per step)
|
||||
- [ ] Zero `cuMemAlloc` calls in `sample_proportional` after first warmup call (pre-allocated buffers only)
|
||||
- [ ] `scripts/gpu-hotpath-guard.sh` catches any new CPU roundtrip in hot path
|
||||
- [ ] H100 epoch time < 5s with 300 training steps, 51 atoms, batch_size ≥ 512
|
||||
- [ ] RTX 3050 epoch time < 10s (proportional improvement from current ~15s)
|
||||
- [ ] All 1861+ tests pass
|
||||
- [ ] Training metrics regression < 5%: eval Sharpe within 5% of baseline, loss convergence within 10% of baseline epoch count
|
||||
- [ ] `to_host()` deprecated with `#[deprecated]` attribute — only `to_host_checkpoint()` available for cold paths
|
||||
|
||||
## Risk: Experience Collector Becomes Bottleneck
|
||||
|
||||
At 2.3% of current epoch time, the experience collector seems negligible. But on H100 with 256 episodes × 500 timesteps = 128K forward passes, it likely takes 5-10s. Once training drops to <5s, experience collection becomes the dominant phase.
|
||||
|
||||
Mitigations (included in Phase 3):
|
||||
- **Phase 3a (dual-stream)**: overlap experience with training → hides latency entirely
|
||||
- **Occupancy**: experience kernel uses `episodes_per_block=4-8` with warp kernel on H100 — already reasonably optimized
|
||||
- **Reduce timesteps**: 500 timesteps may be excessive; 200 could suffice with higher episode count for same total experiences
|
||||
- **Profile independently**: add per-kernel CUDA event timing inside experience collector to identify sub-kernel bottlenecks
|
||||
|
||||
## Non-Goals
|
||||
|
||||
- Multi-GPU training (separate effort, needs NCCL)
|
||||
- Model architecture changes (network dims stay 256×256)
|
||||
- Hyperopt speed (separate binary, different bottleneck)
|
||||
Reference in New Issue
Block a user