diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 066fa7d82..182c14c82 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -869,12 +869,26 @@ impl GpuDqnTrainer { .map_err(|e| MLError::ModelError(format!("stream sync before readback: {e}")))?; // ── 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]; - self.stream.memcpy_dtoh(&self.total_loss_buf, &mut loss_host) - .map_err(|e| MLError::ModelError(format!("DtoH total_loss: {e}")))?; let mut norm_host = [0.0_f32; 1]; - self.stream.memcpy_dtoh(&self.grad_norm_buf, &mut norm_host) - .map_err(|e| MLError::ModelError(format!("DtoH grad_norm: {e}")))?; + 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:?}"))); + } self.scalars_readback_host = [loss_host[0], norm_host[0]]; Ok(FusedTrainScalars {