gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

perf(rl): remove sync_training_event + 24 K-loop memsets

Browse Source 
nsys profiling identified two hot-path bottlenecks:

1. sync_training_event() — cuEventSynchronize blocked host 5-11ms
   EVERY step (87.9% of CUDA API wall time). Removed: LR controller
   uses EMA smoothing, 2-step deferred mapped-pinned reads are fine.

2. 24 raw_memset_d8_zero calls per K-loop iteration — zeroed scratch
   buffers that backward kernels overwrite completely. At K_max=4:
   72 memsets/step, 26% of CUDA API time. Disabled via if-false gate.

Expected combined impact: ~2s saved per 200 steps → measurable sps
improvement at b=1024.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-27 22:44:01 +02:00
parent 2bdb55cc5b
commit 90f178ae9e
1 changed files with 13 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
crates/ml-alpha/src/trainer/integrated.rs
View File

@@ -3782,12 +3782,11 @@ impl IntegratedTrainer {
// The perception trainer's `step_batched_from_device` records a
// CUDA event after its training graph launch instead of
// host-blocking via `cuStreamSynchronize`. By the time we
// reach this point, the GPU has executed the entire reward /
// PER / controller pipeline since the event was recorded, so
// the sync completes instantly (~0 us). This guarantees the
// mapped-pinned loss values below are host-visible.
self.perception.sync_training_event()
.context("sync perception training event")?;
// REMOVED: sync_training_event() blocked the host for 5-11ms
// every step (87.9% of CUDA API wall time per nsys). The LR
// controller uses EMA smoothing over hundreds of steps, so
// 2-step deferred mapped-pinned reads are acceptable. The
// event is still RECORDED for checkpoint save correctness.
// ── Step 0b: read PREVIOUS step's losses from mapped-pinned ────
// Mapped-pinned buffers are GPU+host visible. The GPU wrote loss
@@ -4708,9 +4707,14 @@ impl IntegratedTrainer {
debug_assert_eq!(self.sampled_h_t_d.len(), b_size * HIDDEN_DIM);
debug_assert_eq!(self.sampled_h_tp1_d.len(), b_size * HIDDEN_DIM);
// Zero persistent per-iter scratch — raw memset bypasses cudarc's
// bind_to_thread + event tracking (~200us overhead per call).
{
// REMOVED: 24 raw_memset_d8_zero calls that zeroed DQN/outcome/IQN
// scratch buffers before each K-loop replay iteration. These buffers
// are OVERWRITTEN by backward kernels and reduce_axis0 — the zeroing
// was pure waste. At K_max=4, this was 72 memsets/step. nsys measured
// 458ms total for cuMemsetD8Async (26% of CUDA API time).
// Same pattern already fixed in step_synthetic_body (documented as
// "36 memset_zeros REMOVED").
if false {
let s = self.raw_stream;
unsafe {
raw_memset_d8_zero(self.ss_q_loss_dev_ptr, std::mem::size_of::<f32>(), s)