72 lines
4.0 KiB
Markdown
72 lines
4.0 KiB
Markdown
# Dead Code Cleanup + GPU Per-Phase Profiling
|
|
|
|
## Part 1: Dead Code Cleanup
|
|
|
|
### Problem
|
|
Three replay buffer implementations exist: CPU uniform, CPU PER (segment tree), GPU PER (prefix scan). Only GPU PER is used. The CPU paths are dead code that confuses the architecture and wastes the `ReplayBufferType` enum indirection.
|
|
|
|
### Delete
|
|
- `crates/ml-dqn/src/prioritized_replay.rs` — CPU segment tree PER
|
|
- `crates/ml-dqn/src/prioritized_replay_staleness.rs` — CPU staleness tracker
|
|
- `crates/ml-dqn/src/replay_buffer.rs` — CPU uniform replay buffer
|
|
- `crates/ml-dqn/src/hindsight_replay.rs` — CPU HER (GPU HER is `gpu_her.rs`)
|
|
|
|
### Simplify
|
|
- `crates/ml-dqn/src/replay_buffer_type.rs` — delete `ReplayBufferType` enum. Agent holds `GpuReplayBuffer` directly.
|
|
- `crates/ml-dqn/src/dqn.rs` — `memory` field becomes `GpuReplayBuffer` (was `ReplayBufferType`). `memory()` returns `&GpuReplayBuffer`. `sample()` returns `GpuBatchPtrs` directly.
|
|
- `crates/ml-dqn/src/lib.rs` — remove dead re-exports (`PrioritizedReplayBuffer`, `ReplayBuffer`, `ReplayBufferConfig`, etc.)
|
|
- `crates/ml/src/trainers/dqn/trainer/training_loop.rs` — `agent.memory().sample()` returns `GpuBatchPtrs` directly. Remove `BatchSample` wrapper with empty `experiences` vec. Remove `.gpu_batch.as_ref().unwrap()`.
|
|
- `crates/ml/src/trainers/dqn/fused_training.rs` — `run_full_step` takes `&GpuBatchPtrs` directly (not `&BatchSample`).
|
|
- `crates/ml/src/trainers/dqn/trainer/constructor.rs` — remove `PrioritizedReplayConfig` import used for CPU HER.
|
|
- Tests referencing CPU PER — update or delete.
|
|
|
|
### Non-goals
|
|
- Do not change `gpu_replay_buffer.rs` internals.
|
|
- Do not change `gpu_her.rs` (GPU HER stays).
|
|
- Do not change PER kernel code (`per_kernels.cu`).
|
|
|
|
---
|
|
|
|
## Part 2: GPU Per-Phase CUDA Event Profiling
|
|
|
|
### Problem
|
|
H100 shows 329ms/batch for a 336K param DQN with batch=16384. Expected ~30-40ms. CPU-side wall-clock timers conflate GPU pipeline stalls with actual kernel work. We need GPU-side per-phase timing to identify the real bottleneck.
|
|
|
|
### Design
|
|
Pre-allocate CUDA event pairs in `FusedTrainingCtx`. Record `cuEventRecord` before and after each training phase. At epoch end, compute `cuEventElapsedTime` for each phase and log averages.
|
|
|
|
### Phases to measure
|
|
| Phase | What it covers |
|
|
|-------|---------------|
|
|
| `per_scan` | `per_prefix_scan` kernel (prefix sum over 24.7M priorities) |
|
|
| `per_sample` | `per_sample` + `i64_to_u32` + `gather_*` + `is_weights` kernels |
|
|
| `upload` | `upload_batch_gpu` (DtoD copy from replay to trainer buffers) |
|
|
| `fwd_bwd` | `graph_mega` or `graph_forward` + backward replay |
|
|
| `adam` | `graph_adam` replay (grad norm + optimizer + unflatten) |
|
|
| `per_update` | `per_update_pa` kernel (priority scatter) |
|
|
|
|
### Implementation
|
|
- 12 CUDA events (2 per phase) allocated at `FusedTrainingCtx` construction
|
|
- `cuEventRecord(start, stream)` before each phase
|
|
- `cuEventRecord(end, stream)` after each phase
|
|
- Per-batch: accumulate elapsed times into f64 accumulators
|
|
- Per-epoch: divide by batch count, log one summary line
|
|
- Zero GPU overhead — events are hardware timestamp queries, not sync points
|
|
|
|
### Output format
|
|
```
|
|
GPU phase timing (178 batches avg): per_scan=1.8ms per_sample=3.2ms upload=0.9ms fwd_bwd=22.1ms adam=1.8ms per_update=0.3ms total=30.1ms
|
|
```
|
|
|
|
### Constraints
|
|
- Always-on (no CLI flag). CUDA events have zero overhead.
|
|
- Events must be on the SAME stream as the kernels they bracket.
|
|
- `cuEventElapsedTime` requires both events to have completed — call at epoch end after `cuStreamSynchronize`.
|
|
- Do not add `cuStreamSynchronize` between phases — that would serialize the pipeline and change the behavior we're measuring.
|
|
|
|
### Success criteria
|
|
The timing output reveals where the 329ms goes:
|
|
- If `fwd_bwd` is ~25ms and `total` is ~30ms → the 329ms wall-clock is a pipeline stall (cross-stream sync, cudarc overhead, or CPU scheduling)
|
|
- If `fwd_bwd` is ~280ms → actual GPU kernel bottleneck (occupancy, register pressure, or memory bandwidth)
|
|
- Either answer tells us exactly what to fix next
|