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:
@@ -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.
|
||||
|
||||
@@ -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")?;
|
||||
|
||||
Reference in New Issue
Block a user