diff --git a/crates/ml-core/src/cuda_autograd/gpu_tensor.rs b/crates/ml-core/src/cuda_autograd/gpu_tensor.rs index 44e2f325b..ba0e7506b 100644 --- a/crates/ml-core/src/cuda_autograd/gpu_tensor.rs +++ b/crates/ml-core/src/cuda_autograd/gpu_tensor.rs @@ -277,7 +277,7 @@ impl GpuTensor { let src_start = o * src_stride; let dst_start = o * dst_stride + col_offset; let n = src_stride; - if n > 0 && src_start + n <= t.numel() { + if n > 0 && src_start + n <= t.data.len() && dst_start + n <= output.len() { let src_view = t.data.slice(src_start..src_start + n); let mut dst_view = output.slice_mut(dst_start..dst_start + n); stream.memcpy_dtod(&src_view, &mut dst_view).map_err(|e| { diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index a4447dc6f..33c89e4de 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -785,16 +785,12 @@ impl GpuDqnTrainer { // ── Execute training step ── // Use CUDA Graph when shmem fits in 48KB (graph replay doesn't reliably // inherit cuFuncSetAttribute opt-in for > 48KB on some drivers). - if should_use_cuda_graph(self.shmem_bytes) { - if let Some(ref graph) = self.training_graph { - graph.0.launch().map_err(|e| { - MLError::ModelError(format!("CUDA graph replay: {e}")) - })?; - } else { - self.capture_training_graph(online_dueling, online_branching)?; - } + if let Some(ref graph) = self.training_graph { + graph.0.launch().map_err(|e| { + MLError::ModelError(format!("CUDA graph replay: {e}")) + })?; } else { - self.submit_training_ops(online_dueling, online_branching)?; + self.capture_training_graph(online_dueling, online_branching)?; } // ── Consolidate readback: gather 3 GPU buffers → 1 readback buf ── @@ -852,33 +848,30 @@ impl GpuDqnTrainer { // ── Execute training step ── // Use CUDA Graph when shmem fits in 48KB (graph replay doesn't reliably // inherit cuFuncSetAttribute opt-in for > 48KB on some drivers). - if should_use_cuda_graph(self.shmem_bytes) { - if let Some(ref graph) = self.training_graph { - graph.0.launch().map_err(|e| { - MLError::ModelError(format!("CUDA graph replay: {e}")) - })?; - } else { - self.capture_training_graph(online_dueling, online_branching)?; - } + if let Some(ref graph) = self.training_graph { + graph.0.launch().map_err(|e| { + MLError::ModelError(format!("CUDA graph replay: {e}")) + })?; } else { - self.submit_training_ops(online_dueling, online_branching)?; + self.capture_training_graph(online_dueling, online_branching)?; } - // ── Synchronize + readback (raw — bypasses stale events from graph capture) ── + // ── Raw sync + readback (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!("sync before readback: {r:?}"))); - } + cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); } 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); 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) + cudarc::driver::sys::cuMemcpyDtoH_v2( + loss_host.as_mut_ptr().cast(), + raw_device_ptr(&self.total_loss_buf, &self.stream), 4, + ); + cudarc::driver::sys::cuMemcpyDtoH_v2( + norm_host.as_mut_ptr().cast(), + raw_device_ptr(&self.grad_norm_buf, &self.stream), 4, + ); + } // gpu-exit: 2 scalar readbacks (8 bytes) self.scalars_readback_host = [loss_host[0], norm_host[0]]; Ok(FusedTrainScalars { @@ -1845,13 +1838,11 @@ impl GpuDqnTrainer { ) -> Result<(), MLError> { let sizes = compute_param_sizes(&self.config); - // 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 { - return Err(MLError::ModelError(format!("EMA pre-sync: {r:?}"))); - } - } + // Sync stream and clear any stale errors from graph capture phase. + // cudarc stores errors from cuStreamWaitEvent on disabled events during + // graph capture. check_err() consumes them so bind_to_thread() succeeds. + unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); } + let _ = self.stream.context().check_err(); // Paired (target, online) slices in GOFF_* order (20 pairs) let pairs: [(&CudaSlice, &CudaSlice); 20] = [ @@ -1889,8 +1880,8 @@ impl GpuDqnTrainer { shared_mem_bytes: 0, }; - // Use raw device pointers to bypass stale event tracking on - // graph-modified CudaSlices (events from capture are invalid). + // Use raw device pointers — graph capture left stale events on + // weight CudaSlices, making cudarc's launch_builder.arg() fail. let t_ptr = raw_device_ptr(target_slice, &self.stream); let o_ptr = raw_device_ptr(online_slice, &self.stream); unsafe { @@ -2200,18 +2191,6 @@ pub(crate) fn query_max_shmem_bytes() -> usize { /// Enable only on H100/A100 where the full pipeline has been validated. /// Consumer GPUs use direct kernel launches (slightly higher dispatch overhead /// but avoids the stale event problem). -pub(crate) fn should_use_cuda_graph(_shmem_bytes: usize) -> bool { - use cudarc::driver::sys::{cuDeviceGetAttribute, CUdevice_attribute}; - let mut max_shmem: i32 = 0; - let _ = unsafe { - cuDeviceGetAttribute( - &mut max_shmem, - CUdevice_attribute::CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN, - 0, - ) - }; - max_shmem >= 164_000 -} fn compute_shmem_bytes(config: &GpuDqnTrainConfig) -> usize { let shmem_max_in_dim = config