From 2196cdc60c1b4c669d3ee6ab3bf705a62bf10cdd Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 17 Mar 2026 12:23:34 +0100 Subject: [PATCH] fix(cuda): disable event tracking at stream fork, not at trainer init Moving disable_event_tracking() from GpuDqnTrainer::new() to the stream fork in constructor.rs ensures ALL CudaSlice allocations (experience collector, replay buffer, curiosity trainer) are created with tracking already disabled. Previous placement caused hangs because earlier allocations had stale events. Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 10 +++------- crates/ml/src/trainers/dqn/trainer/constructor.rs | 8 +++++++- 2 files changed, 10 insertions(+), 8 deletions(-) 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