From 68ef8eea3e52a55b7333fcb27c27612e9008b952 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 19 Apr 2026 09:09:46 +0200 Subject: [PATCH] fix: remove ungraphed vaccine batch sampling from training loop MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Vaccine now uses prev_grad_buf inside maintenance_child (graphed). The external sample() call was launching ~10 ungraphed PER kernels on the training stream every 5th step — potential CUmodule corruption and CPU-GPU serialization point. Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/ml/src/trainers/dqn/fused_training.rs | 2 ++ crates/ml/src/trainers/dqn/trainer/training_loop.rs | 13 ++----------- 2 files changed, 4 insertions(+), 11 deletions(-) diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index ec42fbd5d..ed00f8d34 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -909,7 +909,9 @@ impl FusedTrainingCtx { // ── Graph launch or ungraphed step 0 ────────────────────────── if let Some(ref pg) = self.parent_graph { // Steps 1+: single parent graph launch replays all children. + tracing::debug!("parent graph launch START"); pg.launch(cu_stream)?; + tracing::debug!("parent graph launch DONE (async — GPU working)"); } else { // Ungraphed step 0: run everything, then capture + compose. // Ensure params are initialized for the first step. diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 743540ee6..9cd892fb1 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -1514,22 +1514,13 @@ impl DQNTrainer { break; } } - 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() - } else { - None - }; - sample_total_us += sample_start.elapsed().as_micros() as u64; + // Vaccine batch sampling removed — vaccine now uses prev_grad_buf + // inside maintenance_child (graphed). No ungraphed PER sampling needed. { let fused_start = std::time::Instant::now(); let fused = self.fused_ctx.as_mut() .expect("Fused CUDA training is the only production path"); - if let Some(vb) = vaccine_batch { - fused.pending_vaccine_batch = Some(vb); - } // 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")?;