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

fix: bypass stale cudarc events for all post-graph operations

Browse Source 
After CUDA Graph capture (which disables/re-enables event tracking),
CudaSlices modified during capture have stale write events. Any cudarc
API call (memcpy_dtoh, synchronize, launch_builder.arg) on these slices
fails with CUDA_ERROR_INVALID_VALUE because cuStreamWaitEvent gets an
invalid event handle.

Fix: use raw CUDA driver calls for ALL post-graph operations:
- cuStreamSynchronize instead of cudarc synchronize() (avoids bind_to_thread)
- cuMemcpyDtoH_v2 via raw_device_ptr for scalar readbacks (2 × 4 bytes)
- raw_device_ptr for EMA kernel args (bypasses device_ptr event wait)

Event tracking remains enabled globally — only the 3 post-graph call sites
use raw driver calls. All other cudarc operations (cat/stack, alloc, etc.)
work normally with event tracking.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-20 23:25:58 +01:00
parent 305b19f6f8
commit 4a06e99ee5
1 changed files with 16 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -864,37 +864,21 @@ impl GpuDqnTrainer {
self.submit_training_ops(online_dueling, online_branching)?;
}
// ── Synchronize stream before readback ──────────────────────
// Raw sync: cudarc's synchronize() propagates stale errors from
// graph capture via bind_to_thread/check_err. Raw is clean.
// ── Synchronize + readback (raw — bypasses stale events from graph capture) ──
unsafe {
let r = cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream());
if r != cudarc::driver::sys::CUresult::CUDA_SUCCESS {
return Err(MLError::ModelError(format!("stream sync before readback: {r:?}")));
return Err(MLError::ModelError(format!("sync before readback: {r:?}")));
}
}
// ── Gather 2 scalars to host: total_loss + grad_norm ────────
// These buffers were written by the CUDA Graph (which runs with event
// tracking disabled). Their CudaSlice write events are stale, so
// cudarc's memcpy_dtoh would fail with CUDA_ERROR_INVALID_VALUE.
// Use raw_device_ptr (ManuallyDrop guard) + synchronous cuMemcpyDtoH.
let mut loss_host = [0.0_f32; 1];
let mut norm_host = [0.0_f32; 1];
let loss_ptr = raw_device_ptr(&self.total_loss_buf, &self.stream);
let norm_ptr = raw_device_ptr(&self.grad_norm_buf, &self.stream);
let r1 = unsafe { cudarc::driver::sys::cuMemcpyDtoH_v2(
loss_host.as_mut_ptr().cast(), loss_ptr, 4,
)};
if r1 != cudarc::driver::sys::CUresult::CUDA_SUCCESS {
return Err(MLError::ModelError(format!("DtoH total_loss: {r1:?}")));
}
let r2 = unsafe { cudarc::driver::sys::cuMemcpyDtoH_v2(
norm_host.as_mut_ptr().cast(), norm_ptr, 4,
)};
if r2 != cudarc::driver::sys::CUresult::CUDA_SUCCESS {
return Err(MLError::ModelError(format!("DtoH grad_norm: {r2:?}")));
}
unsafe {
cudarc::driver::sys::cuMemcpyDtoH_v2(loss_host.as_mut_ptr().cast(), loss_ptr, 4);
cudarc::driver::sys::cuMemcpyDtoH_v2(norm_host.as_mut_ptr().cast(), norm_ptr, 4);
} // gpu-exit: 2 scalar readbacks (8 bytes total)
self.scalars_readback_host = [loss_host[0], norm_host[0]];
Ok(FusedTrainScalars {
@@ -1203,8 +1187,8 @@ impl GpuDqnTrainer {
self.stream.synchronize()
.map_err(|e| MLError::ModelError(format!("stream sync before capture: {e}")))?;
// Disable event tracking during capture — cudarc's device_ptr/device_ptr_mut
// record CudaEvents which are DISALLOWED inside CUDA Graph capture.
// Disable event tracking during capture — cudarc's device_ptr records
// CudaEvents which are DISALLOWED inside CUDA Graph capture.
unsafe { self.stream.context().disable_event_tracking(); }
// Begin stream capture — only work submitted from this thread on this
@@ -1226,9 +1210,7 @@ impl GpuDqnTrainer {
cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
);
// Re-enable event tracking after capture. With GPU-native cat/stack
// (no to_host roundtrip), stale events from capture don't cause issues
// because all post-capture GPU ops use DtoD copies, not DtoH.
// Re-enable event tracking after capture.
unsafe { self.stream.context().enable_event_tracking(); }
// Propagate submission error first
@@ -1863,8 +1845,7 @@ impl GpuDqnTrainer {
) -> Result<(), MLError> {
let sizes = compute_param_sizes(&self.config);
// Raw sync: cudarc's synchronize() may propagate stale errors from
// graph capture via check_err(). Raw cuStreamSynchronize is clean.
// Raw sync — bypasses cudarc's bind_to_thread which propagates stale errors.
unsafe {
let r = cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream());
if r != cudarc::driver::sys::CUresult::CUDA_SUCCESS {
@@ -1908,15 +1889,15 @@ impl GpuDqnTrainer {
shared_mem_bytes: 0,
};
// Safety: argument order matches dqn_ema_kernel(float* target, const float* online, float tau, int n).
// target and online CudaSlice<f32> buffers have size >= sizes[i].
// Both allocated on the same forked stream. Stream synchronized above
// to ensure CUDA Graph execution is complete.
// Use raw device pointers to bypass stale event tracking on
// graph-modified CudaSlices (events from capture are invalid).
let t_ptr = raw_device_ptr(target_slice, &self.stream);
let o_ptr = raw_device_ptr(online_slice, &self.stream);
unsafe {
self.stream
.launch_builder(&self.ema_kernel)
.arg(*target_slice)
.arg(*online_slice)
.arg(&t_ptr)
.arg(&o_ptr)
.arg(&tau)
.arg(&n)
.launch(launch_cfg)