diff --git a/crates/ml-alpha/src/trainer/integrated.rs b/crates/ml-alpha/src/trainer/integrated.rs index bd27056cb..25268eb40 100644 --- a/crates/ml-alpha/src/trainer/integrated.rs +++ b/crates/ml-alpha/src/trainer/integrated.rs @@ -427,6 +427,8 @@ pub struct IntegratedTrainer { postfill_graph: Option, reward_graph: Option, graph_warmup_done: bool, + training_graph: Option, + training_warmup_done: bool, // ── Kernel handles for grad combine + cross-batch reduce ────────── _grad_h_module: Arc, @@ -1905,6 +1907,8 @@ impl IntegratedTrainer { postfill_graph: None, reward_graph: None, graph_warmup_done: false, + training_graph: None, + training_warmup_done: false, _grad_h_module: grad_h_module, grad_h_accumulate_fn, _reduce_axis0_module: reduce_axis0_module, @@ -3315,6 +3319,28 @@ impl IntegratedTrainer { let h_t_borrow: &CudaSlice = self.perception.h_t_view(); debug_assert_eq!(h_t_borrow.len(), b_size * HIDDEN_DIM); + // ── CUDA Graph C: replay training step ─────────────────────── + // Three-state machine: first step = warmup (eager), second = + // capture, third+ = replay. All device pointers are stable + // (pre-allocated ss_* fields). Stream sync + loss readback + // stay OUTSIDE the captured region. + if self.training_graph.is_some() { + self.training_graph + .as_ref() + .unwrap() + .launch() + .context("training graph launch")?; + } else { + let capturing_training = self.training_warmup_done; + if !self.training_warmup_done { + self.training_warmup_done = true; + } + if capturing_training { + self.stream + .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED) + .map_err(|e| anyhow::anyhow!("training begin_capture: {e}"))?; + } + // ── Step 3: zero persistent per-step scratch ───────────────── // Buffers are pre-allocated struct fields for CUDA Graph pointer // stability. memset_zeros resets them each call. @@ -3457,9 +3483,7 @@ impl IntegratedTrainer { &mut self.ss_pi_loss_entropy_d, ) .context("policy_head.surrogate_forward")?; - let l_pi_host = read_scalar_d(&self.stream, &self.ss_pi_loss_d)?; - let _l_pi_ent_host = read_scalar_d(&self.stream, &self.ss_pi_loss_entropy_d)?; - self.last_pi_loss = l_pi_host; + // Host reads deferred to after graph capture region. } // ── Step 6a: Bellman target build (item 2) ─────────────────── @@ -3508,9 +3532,7 @@ impl IntegratedTrainer { &mut self.ss_q_grad_logits_d, ) .context("dqn_head.backward_logits (sampled_actions) [R7d off-policy]")?; - let l_q_host = read_scalar_d(&self.stream, &self.ss_q_loss_d)?; - // Mirror for next step's LR controller plateau detection. - self.last_q_loss = l_q_host / (b_size as f32); + // Host read deferred to after graph capture region. self.dqn_head .backward_to_w_b_h( @@ -3615,10 +3637,7 @@ impl IntegratedTrainer { reduce_axis0_free(&self.stream, &self.reduce_axis0_fn, &self.ss_v_grad_w_per_batch_d, b_size, HIDDEN_DIM, &mut self.ss_v_grad_w_d)?; reduce_axis0_free(&self.stream, &self.reduce_axis0_fn, &self.ss_v_grad_b_per_batch_d, b_size, 1, &mut self.ss_v_grad_b_d)?; reduce_axis0_free(&self.stream, &self.reduce_axis0_fn, &self.ss_v_loss_per_batch_d, b_size, 1, &mut self.ss_v_loss_sum_d)?; - let l_v_sum_host = read_scalar_d(&self.stream, &self.ss_v_loss_sum_d)?; - let l_v_host = l_v_sum_host / (b_size as f32); - // Mirror for next step's LR controller plateau detection. - self.last_v_loss = l_v_host; + // Host read deferred to after graph capture region. // ── Step 9: Adam updates on each head's w and b ────────────── self.dqn_w_adam @@ -3649,8 +3668,7 @@ impl IntegratedTrainer { // Adam steps are no-ops modulo β decay. Once F.5 lands real // labels, this path activates the supervised learning signal // without any further trainer-side changes. - let l_frd_host: f32 = { - let frd_n_h = crate::rl::common::FRD_N_HORIZONS; + { let frd_h_dim = crate::rl::common::FRD_HIDDEN_DIM; let frd_out_dim = crate::rl::frd::FRD_OUT_DIM; // Zero FRD scratch buffers. @@ -3719,11 +3737,8 @@ impl IntegratedTrainer { self.frd_b2_adam .step(&mut self.frd_head.b2_d, &self.ss_frd_grad_b2_d) .context("frd_b2_adam.step")?; - // Per-row CE loss → mean over (b, h). Mapped-pinned read. - let loss_pb_h = - read_slice_d(&self.stream, &self.ss_frd_loss_per_b_h_d, b_size * frd_n_h)?; - loss_pb_h.iter().sum::() / ((b_size * frd_n_h) as f32) - }; + // Host read deferred to after graph capture region. + } // ── Step 10: encoder grad combine (Phase E.2-DEFER item 1) ─── // Zero the combined slot, fold π+V grad_h_t with λ-weights. @@ -3917,6 +3932,37 @@ impl IntegratedTrainer { ) .context("ema_update_per_step(v_grad_norm)")?; + if capturing_training { + let graph = self.stream + .end_capture( + CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH, + ) + .context("training end_capture")? + .ok_or_else(|| anyhow::anyhow!("training end_capture returned None"))?; + self.training_graph = Some(graph); + } + } // end else (training warmup / capture) + + // ── Shared path: sync + deferred loss readback ─────────────── + // Both the graph-launch path and the warmup/capture path need + // to synchronize before reading mapped-pinned loss scalars. + self.stream.synchronize().context("step_synthetic sync")?; + + let l_pi_host = read_scalar_d(&self.stream, &self.ss_pi_loss_d)?; + self.last_pi_loss = l_pi_host; + + let l_q_host = read_scalar_d(&self.stream, &self.ss_q_loss_d)?; + self.last_q_loss = l_q_host / (b_size as f32); + + let l_v_sum_host = read_scalar_d(&self.stream, &self.ss_v_loss_sum_d)?; + let l_v_host = l_v_sum_host / (b_size as f32); + self.last_v_loss = l_v_host; + + let frd_n_h = crate::rl::common::FRD_N_HORIZONS; + let loss_pb_h = + read_slice_d(&self.stream, &self.ss_frd_loss_per_b_h_d, b_size * frd_n_h)?; + let l_frd_host = loss_pb_h.iter().sum::() / ((b_size * frd_n_h) as f32); + // ── Step 12: compose stats ─────────────────────────────────── // BCE / aux losses are NOT read this phase — perception is driven // separately by callers via its existing step_batched path. The