diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 9cabc4c52..a4447dc6f 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -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 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)