diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 7339d02b8..c22ec6bbc 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -400,13 +400,9 @@ impl GpuDqnTrainer { let b = config.batch_size; let total_params = compute_total_params(&config); - // All GPU components share this single forked CudaStream. Disable cudarc's - // automatic event tracking: it injects cuStreamWaitEvent/cuEventRecord into - // every CudaSlice access, creating cross-stream event references that - // invalidate CUDA Graph capture and conflict with Candle tensor Drops on - // the default stream. Safe: single forked stream, all work serialized. - // SAFETY: single-stream, all GPU work is serialized on `stream`. - unsafe { stream.context().disable_event_tracking(); } + // Event tracking already disabled at stream fork time (constructor.rs). + // All CudaSlice allocations on this stream have no events — safe for + // CUDA Graph capture and single-stream serialized execution. // ── Compile all 6 training kernels from same module ────────── let (forward_loss_kernel, backward_kernel, grad_norm_kernel, adam_update_kernel, f32_to_bf16_kernel, bf16_to_f32_kernel) = diff --git a/crates/ml/src/trainers/dqn/trainer/constructor.rs b/crates/ml/src/trainers/dqn/trainer/constructor.rs index 2a4365526..d302e4915 100644 --- a/crates/ml/src/trainers/dqn/trainer/constructor.rs +++ b/crates/ml/src/trainers/dqn/trainer/constructor.rs @@ -136,7 +136,13 @@ impl DQNTrainer { if let candle_core::Device::Cuda(ref cuda_dev) = device { let stream = cuda_dev.cuda_stream().fork() .map_err(|e| anyhow::anyhow!("Failed to fork CUDA stream: {e}"))?; - info!("Forked dedicated CudaStream for all GPU components"); + // Disable cudarc event tracking BEFORE any CudaSlice allocation. + // Must happen here — if delayed to GpuDqnTrainer::new(), earlier + // CudaSlice allocations (experience collector, replay buffer) will + // have stale events that cause hangs when tracking is later disabled. + // Safe: single forked stream, all GPU work is serialized. + unsafe { stream.context().disable_event_tracking(); } + info!("Forked dedicated CudaStream for all GPU components (event tracking disabled)"); Some(stream) } else { None