From 1d3a533f9933c4f8a5ece6df1a19d9ac7649a088 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 19 Apr 2026 01:14:53 +0200 Subject: [PATCH] =?UTF-8?q?fix:=20capture=20PER=20sampling=20inside=20pare?= =?UTF-8?q?nt=20graph=20=E2=80=94=20zero=20ungraphed=20PER=20kernels=20per?= =?UTF-8?q?=20step?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Move sample_proportional (prefix scan, binary search, gather x6, IS weights) from training_loop.rs into run_full_step, captured as per_sample child graph node. Steps 1+ replay all 13 children via single cuGraphLaunch with no ungraphed PER kernel dispatches. Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/ml/src/trainers/dqn/fused_training.rs | 41 +++++++++++++++---- .../src/trainers/dqn/trainer/training_loop.rs | 7 ++-- 2 files changed, 36 insertions(+), 12 deletions(-) diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 6b054c6b7..22981ba44 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -898,7 +898,6 @@ impl FusedTrainingCtx { /// After graph launch: only host-side pinned scalar reads (nanoseconds, no GPU ops). pub(crate) fn run_full_step( &mut self, - gpu_batch: &crate::dqn::replay_buffer_type::GpuBatch, agent: &mut DQNAgentType, _device: &MlDevice, ) -> Result { @@ -928,6 +927,21 @@ impl FusedTrainingCtx { self.trainer.target_params_initialized = true; } + // Phase -1: PER sampling (writes directly to trainer buffers via gather_padded) + let ptrs = agent.primary_dqn_mut().memory.gpu.sample_proportional(self.batch_size) + .map_err(|e| anyhow::anyhow!("PER sample step 0: {e}"))?; + let gpu_batch = crate::dqn::replay_buffer_type::GpuBatch { + states_ptr: ptrs.states_ptr, + actions_ptr: ptrs.actions_ptr, + rewards_ptr: ptrs.rewards_ptr, + next_states_ptr: ptrs.next_states_ptr, + dones_ptr: ptrs.dones_ptr, + weights_ptr: ptrs.weights_ptr, + indices_ptr: ptrs.indices_ptr, + episode_ids_ptr: ptrs.episode_ids_ptr, + batch_size: self.batch_size, + }; + // Phase 0: Counter increments + stochastic depth mask (ungraphed) self.submit_counters_ops()?; @@ -965,7 +979,7 @@ impl FusedTrainingCtx { .map_err(|e| anyhow::anyhow!("ddqn ops: {e}"))?; // Phase 4: Aux ops (HER + EMA + IQL + IQN + attention + CQL) - self.submit_aux_ops(agent, gpu_batch)?; + self.submit_aux_ops(agent, &gpu_batch)?; // Phase 5b: Post-aux ops (selectivity + denoise + Q-stats + risk_sgd) self.trainer.submit_post_aux_ops(self.batch_size) @@ -988,11 +1002,11 @@ impl FusedTrainingCtx { .map_err(|e| anyhow::anyhow!("maintenance ops: {e}"))?; // Phase 8: IQL modulate + PER priority update (ungraphed) - self.submit_iql_modulate_ops(agent, gpu_batch)?; + self.submit_iql_modulate_ops(agent, &gpu_batch)?; self.submit_per_priority_ops(agent)?; // Capture all children + compose parent after ungraphed step succeeds. - if let Err(e) = self.capture_training_graph(agent, gpu_batch) { + if let Err(e) = self.capture_training_graph(agent, &gpu_batch) { tracing::warn!("parent graph capture failed (continuing ungraphed): {e}"); } } @@ -1519,13 +1533,20 @@ impl FusedTrainingCtx { /// independently, then composed into the parent via `cuGraphAddChildGraphNode`. /// /// Children (sequential): - /// counters -> spectral -> forward -> ddqn -> aux -> post_aux -> + /// per_sample -> counters -> spectral -> forward -> ddqn -> aux -> post_aux -> /// adam_grad -> adam -> maintenance -> iql_modulate -> per_priority fn capture_training_graph( &mut self, agent: &mut DQNAgentType, gpu_batch: &crate::dqn::replay_buffer_type::GpuBatch, ) -> Result<()> { + // Capture PER sampling child (prefix scan + binary search + gather). + let per_sample = self.capture_child_graph("per_sample", |s| { + agent.primary_dqn_mut().memory.gpu.sample_proportional(s.batch_size) + .map_err(|e| anyhow::anyhow!("per_sample capture: {e}")) + .map(|_| ()) + })?; + // Capture counters child (GPU-side step counters + stochastic depth mask). let counters = self.capture_child_graph("counters", |s| { s.submit_counters_ops() @@ -1594,6 +1615,7 @@ impl FusedTrainingCtx { })?; // Store children. + self.per_sample_child = Some(per_sample); self.counters_child = Some(counters); self.spectral_child = Some(spectral); self.forward_child = Some(forward); @@ -1612,14 +1634,14 @@ impl FusedTrainingCtx { // Create eval-only forward exec for deterministic Q-value replay. self.create_eval_forward_exec()?; - info!("12 child graphs captured + parent composed (counters + spectral + forward + ddqn + aux + post_aux + adam_grad + adam + maintenance + iql_modulate + per_priority)"); + info!("13 child graphs captured + parent composed (per_sample + counters + spectral + forward + ddqn + aux + post_aux + adam_grad + adam + maintenance + iql_modulate + per_priority)"); Ok(()) } /// Compose parent graph from all child sub-graphs using `cuGraphAddChildGraphNode`. /// /// Each child is added as a sequential dependency: - /// counters -> spectral -> forward -> ddqn -> aux -> post_aux -> + /// per_sample -> counters -> spectral -> forward -> ddqn -> aux -> post_aux -> /// adam_grad -> adam -> maintenance -> iql_modulate -> per_priority /// The parent is instantiated once with `cuGraphInstantiateWithFlags`. fn compose_parent_graph(&mut self) -> Result<()> { @@ -1632,7 +1654,10 @@ impl FusedTrainingCtx { // Collect all children in sequential order. // Optional children are skipped if None. - let mut children_vec: Vec<(&str, &ChildGraph)> = Vec::with_capacity(12); + let mut children_vec: Vec<(&str, &ChildGraph)> = Vec::with_capacity(13); + if let Some(ref c) = self.per_sample_child { + children_vec.push(("per_sample", c)); + } if let Some(ref c) = self.counters_child { children_vec.push(("counters", c)); } diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index cef8d5da9..dc8671acb 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -1506,15 +1506,13 @@ impl DQNTrainer { let mut agent = self.agent.write().await; for _step in 0..num_training_steps { - let sample_start = std::time::Instant::now(); { let buffer = agent.memory(); if !buffer.can_sample(self.current_batch_size) { break; } } - let batch = agent.memory_mut().sample(self.current_batch_size) - .map_err(|e| anyhow::anyhow!("PER sample: {e}"))?; + let sample_start = std::time::Instant::now(); // Vaccine batch: only sample when needed (every 5th step for better diversity) let vaccine_batch = if train_step_count % 5 == 0 { agent.memory_mut().sample(self.current_batch_size).ok() @@ -1530,7 +1528,8 @@ impl DQNTrainer { if let Some(vb) = vaccine_batch { fused.pending_vaccine_batch = Some(vb); } - let _fused_result = fused.run_full_step(&batch, &mut *agent, &self.device) + // PER sampling is inside run_full_step (per_sample child graph). + let _fused_result = fused.run_full_step(&mut *agent, &self.device) .context("Fused CUDA training step failed")?; let step_us = fused_start.elapsed().as_micros() as u64;