fix: remove ungraphed vaccine batch sampling from training loop

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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-19 09:09:46 +02:00
parent c2edb273ca
commit 68ef8eea3e
2 changed files with 4 additions and 11 deletions

View File

@@ -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.

View File

@@ -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")?;