gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(cuda): disable event tracking at stream fork, not at trainer init

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-17 12:23:34 +01:00
parent b39d10c968
commit 2196cdc60c
2 changed files with 10 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -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) =

8
crates/ml/src/trainers/dqn/trainer/constructor.rs
View File

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