diff --git a/crates/ml-alpha/src/trainer/integrated.rs b/crates/ml-alpha/src/trainer/integrated.rs index 51c1c6f78..1df82bb1d 100644 --- a/crates/ml-alpha/src/trainer/integrated.rs +++ b/crates/ml-alpha/src/trainer/integrated.rs @@ -328,6 +328,14 @@ const RL_HINDSIGHT_FORWARD_CUBIN: &[u8] = /// Phase E.3+ without touching the kernel. const RL_LR_CONTROLLER_ALPHA: f32 = 0.4; +/// Spectral norm power iteration period. Full power iteration (175us per +/// head) runs every SPECTRAL_NORM_PERIOD replay steps; intermediate steps +/// skip entirely. By Weyl's inequality, sigma_max shifts by at most +/// ||Delta_W||_2 per Adam step, so the guard remains valid between +/// full iterations. 100 steps at 175us/head x 3 heads = 843us/step +/// amortized to 8.4us/step. Always runs (not a feature flag). +const SPECTRAL_NORM_PERIOD: u64 = 100; + /// Pre-extracted raw CUdeviceptr values for hot-path buffers. /// All pointers come from pre-allocated `CudaSlice` fields whose /// device addresses are stable for the trainer's lifetime. @@ -479,6 +487,11 @@ pub struct IntegratedTrainer { /// Recorded on `self.stream` after dqn_replay_step (k_iter > 0); /// `train_stream` waits on this before priority update. replay_done_event: CudaEvent, + /// Recorded on `train_stream` after the last kernel in the K-loop + /// (priority tree rebuild). `step_with_lobsim` waits on this event + /// instead of `train_stream.synchronize()` — event wait is GPU-side + /// only (no host stall), saving ~2ms/step host idle time. + train_done_event: CudaEvent, // ── CUDA Graph capture for the RL step pipeline ───────────────── // Three-state machine: first step = warmup (eager), second = @@ -952,6 +965,15 @@ pub struct IntegratedTrainer { /// the same config + step sequence. step_counter: u64, + /// Host-side counter for spectral norm amortization. Spectral norm + /// (power iteration on full weight matrices) costs 175us x 3 heads = + /// 843us/step. Running every step is wasteful because per-step weight + /// changes are O(lr * grad) = small. Amortizing to every + /// SPECTRAL_NORM_PERIOD steps via Weyl's inequality: sigma_max + /// shifts by at most ||DeltaW||_2 per step, so the guard remains + /// valid between full power iterations. + spectral_norm_counter: u64, + /// SP20 P3 Forward-Return-Distribution head (forecast over 3 /// horizons × 21 return-bucket atoms). Forward kernel runs every /// `step_with_lobsim`; backward + label-supervised loss arrive in @@ -1049,8 +1071,10 @@ pub struct IntegratedTrainer { isv_staging: MappedF32Buffer, /// FRD loss readback staging `[B × FRD_N_HORIZONS]`. frd_loss_staging: MappedF32Buffer, - /// Scalar loss readback staging `[1]` — reused across all read_scalar_d - /// call sites within a single step (serialised by stream order). + /// Scalar loss readback staging `[1]` — reused for deferred (one-step + /// delayed) loss reads. dqn_replay_step queues async DtoD into this + /// buffer; the NEXT iteration reads the previous value via + /// volatile read. Serialised by stream order across consecutive calls. scalar_staging: MappedF32Buffer, /// Batched loss readback staging `[3]` — holds pi_loss, q_loss, v_loss /// in a single DtoD + sync instead of 3 separate read_scalar_via_staging @@ -2160,6 +2184,9 @@ impl IntegratedTrainer { let replay_done_event = ctx_for_events .new_event(None) .map_err(|e| anyhow::anyhow!("replay_done_event: {e}"))?; + let train_done_event = ctx_for_events + .new_event(None) + .map_err(|e| anyhow::anyhow!("train_done_event: {e}"))?; let hot = HotPathPtrs { isv: isv_d.raw_ptr(), @@ -2197,6 +2224,7 @@ impl IntegratedTrainer { push_done_event, sample_done_event, replay_done_event, + train_done_event, prefill_graph: None, postfill_graph: None, reward_graph: None, @@ -2413,6 +2441,7 @@ impl IntegratedTrainer { last_v_loss: 0.0, last_k_updates: 0, step_counter: 0, + spectral_norm_counter: 0, frd_head, frd_hidden_d, frd_logits_d, @@ -4890,10 +4919,26 @@ impl IntegratedTrainer { &mut self.ss_q_grad_logits_d, ) .context("dqn_replay_step: dqn_head.backward_logits")?; - let l_q_host = read_scalar_via_staging( - &self.stream, &self.ss_q_loss_d, &self.scalar_staging, - )?; - let l_q = l_q_host / (b_size as f32); + // Deferred loss read: read PREVIOUS iteration's Q loss from the + // staging buffer (populated by the async DtoD queued at the end of + // the prior dqn_replay_step or step_synthetic, flushed by + // step_synthetic's batched-loss sync). Then queue THIS iteration's + // Q loss DtoD for the NEXT read. One-step delay on per-branch LR + // controller's IQN-loss proxy is acceptable (slow EMA). At + // construction / first call the staging holds zero = sentinel, and + // the per-branch LR controller's cold-start gate handles that. + let l_q_host = unsafe { std::ptr::read_volatile(self.scalar_staging.host_ptr) }; + let l_q = l_q_host / (b_size as f32).max(1.0); + // Queue async DtoD for NEXT read — no sync needed. + unsafe { + cudarc::driver::result::memcpy_dtod_async( + self.scalar_staging.dev_ptr, + self.hot.ss_q_loss, + std::mem::size_of::(), + self.stream.cu_stream(), + ) + .context("dqn_replay_step: q_loss staging DtoD (async, for next iter)")?; + } self.dqn_head .backward_to_w_b_h( @@ -4923,9 +4968,14 @@ impl IntegratedTrainer { .context("dqn_replay_step: dqn_b_adam.step")?; // ── 6. Spectral norm on DQN, IQN, policy weights ──────────── - // One power iteration per step per weight matrix. The v_buffer is - // persistent for warm-start. If sigma > ISV[SPECTRAL_NORM_MAX], - // the kernel rescales W in place. + // Amortized to every SPECTRAL_NORM_PERIOD replay steps (Weyl's + // inequality: sigma_max shifts by at most ||Delta_W||_2 per Adam + // step, so the guard remains valid between full power iterations). + // 175us/head x 3 heads = 843us/step → 8.4us amortized at + // SPECTRAL_NORM_PERIOD=100. The v_buffer is persistent for + // warm-start across amortized calls. + self.spectral_norm_counter += 1; + if self.spectral_norm_counter % SPECTRAL_NORM_PERIOD == 0 { { let dqn_rows = (N_ACTIONS * Q_N_ATOMS) as i32; let dqn_cols = HIDDEN_DIM as i32; @@ -4989,6 +5039,7 @@ impl IntegratedTrainer { ).map_err(|e| anyhow::anyhow!("rl_spectral_norm(policy_w): {:?}", e))?; } } + } // end spectral norm amortization guard // ── 7. Q-bias correction ───────────────────────────────────── // Track mean(Q_predicted - actual_return), emit correction to ISV. @@ -5144,12 +5195,16 @@ impl IntegratedTrainer { ); } - // ── Sync train_stream: wait for previous step's PER priority - // update + tree rebuild to complete before this step touches - // the replay buffer (push_ring/push_flush write the same tree). - self.train_stream - .synchronize() - .context("wait previous step training (train_stream)")?; + // ── Event-based inter-stream sync: wait for previous step's PER + // priority update + tree rebuild to complete before this step + // touches the replay buffer (push_ring/push_flush write the same + // tree). Uses train_done_event (GPU-side wait) instead of + // train_stream.synchronize() (host-blocking) — the host thread + // stays free to queue subsequent kernels while the GPU serialises + // the dependency. Saves ~2ms/step host idle time. + self.stream + .wait(&self.train_done_event) + .map_err(|e| anyhow::anyhow!("stream wait train_done: {e}"))?; // ── Step 0: bump device-resident step counter (ISV[548]). // Must run BEFORE any kernel that reads current_step from ISV. @@ -6657,6 +6712,16 @@ impl IntegratedTrainer { } } + // Record train_done_event on train_stream after the last K-loop + // iteration's tree rebuild. The NEXT step_with_lobsim's + // stream.wait(&train_done_event) ensures the replay buffer is + // coherent before push_ring/push_flush write new transitions. + // GPU-side event wait replaces the prior host-blocking + // train_stream.synchronize() — saves ~2ms/step. + self.train_done_event + .record(&self.train_stream) + .map_err(|e| anyhow::anyhow!("train_done_event record: {e}"))?; + // Target-net soft update (Phase R5 + R6) — runs once per step // after the Q-head Adam update inside step_synthetic. Reads τ // from ISV[401]. @@ -7170,31 +7235,6 @@ fn write_slice_i32_d( Ok(()) } -/// Hot-path scalar read reusing a pre-allocated `MappedF32Buffer` staging -/// slot. Same DtoD + sync + volatile-read pattern as `read_scalar_d` but -/// without per-call `cuMemHostAlloc` / `cuMemFreeHost`. The staging buffer -/// is `[1]` and serialised by stream order across consecutive calls. -fn read_scalar_via_staging( - stream: &Arc, - src: &CudaSlice, - staging: &MappedF32Buffer, -) -> Result { - debug_assert!(src.len() >= 1); - debug_assert!(staging.len >= 1); - unsafe { - let (src_ptr, _g) = src.device_ptr(stream); - cudarc::driver::result::memcpy_dtod_async( - staging.dev_ptr, - src_ptr, - std::mem::size_of::(), - stream.cu_stream(), - ) - .context("read_scalar_via_staging DtoD")?; - } - stream.synchronize().context("read_scalar_via_staging sync")?; - Ok(unsafe { std::ptr::read_volatile(staging.host_ptr) }) -} - /// Phase E.3b helper: read `n` floats from a device buffer into a host /// `Vec` via mapped-pinned staging + DtoD copy + stream sync. Used /// by `step_with_lobsim` to pull Q-logits / V / π-logits to host for