4.0 KiB
4.0 KiB
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 PERcrates/ml-dqn/src/prioritized_replay_staleness.rs— CPU staleness trackercrates/ml-dqn/src/replay_buffer.rs— CPU uniform replay buffercrates/ml-dqn/src/hindsight_replay.rs— CPU HER (GPU HER isgpu_her.rs)
Simplify
crates/ml-dqn/src/replay_buffer_type.rs— deleteReplayBufferTypeenum. Agent holdsGpuReplayBufferdirectly.crates/ml-dqn/src/dqn.rs—memoryfield becomesGpuReplayBuffer(wasReplayBufferType).memory()returns&GpuReplayBuffer.sample()returnsGpuBatchPtrsdirectly.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()returnsGpuBatchPtrsdirectly. RemoveBatchSamplewrapper with emptyexperiencesvec. Remove.gpu_batch.as_ref().unwrap().crates/ml/src/trainers/dqn/fused_training.rs—run_full_steptakes&GpuBatchPtrsdirectly (not&BatchSample).crates/ml/src/trainers/dqn/trainer/constructor.rs— removePrioritizedReplayConfigimport used for CPU HER.- Tests referencing CPU PER — update or delete.
Non-goals
- Do not change
gpu_replay_buffer.rsinternals. - 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
FusedTrainingCtxconstruction cuEventRecord(start, stream)before each phasecuEventRecord(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.
cuEventElapsedTimerequires both events to have completed — call at epoch end aftercuStreamSynchronize.- Do not add
cuStreamSynchronizebetween 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_bwdis ~25ms andtotalis ~30ms → the 329ms wall-clock is a pipeline stall (cross-stream sync, cudarc overhead, or CPU scheduling) - If
fwd_bwdis ~280ms → actual GPU kernel bottleneck (occupancy, register pressure, or memory bandwidth) - Either answer tells us exactly what to fix next