diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 5370626b4..1886e9042 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -366,6 +366,10 @@ pub struct GpuDqnTrainer { // Dead kernels (forward_loss_kernel, forward_only_kernel, backward_kernel) // removed — replaced by cuBLAS forward/backward. See Task 2 for method cleanup. grad_norm_kernel: CudaFunction, + /// Separate instance of grad_norm for non-graph launches (clip_grad_buf_inplace). + /// CUDA doesn't allow the same CUfunction to be launched both inside a captured + /// graph and outside on the same stream. + grad_norm_standalone: CudaFunction, adam_update_kernel: CudaFunction, ema_kernel: CudaFunction, per_update_kernel: CudaFunction, @@ -1320,7 +1324,7 @@ impl GpuDqnTrainer { unsafe { self.stream - .launch_builder(&self.grad_norm_kernel) + .launch_builder(&self.grad_norm_standalone) .arg(&scratch_ptr) .arg(&norm_ptr) .arg(&total) @@ -1517,10 +1521,12 @@ impl GpuDqnTrainer { /// from overwhelming the subsequent auxiliary gradient additions. /// All operations are async on the same stream — zero CPU sync. pub fn clip_grad_buf_inplace(&mut self, max_norm: f32) -> Result<(), MLError> { + // clip_grad_buf_inplace: standalone kernel launch (not captured in CUDA graph) // Zero the grad_norm accumulator self.stream .memset_zeros(&mut self.grad_norm_buf) .map_err(|e| MLError::ModelError(format!("zero grad_norm for clip: {e}")))?; + // memset_zeros succeeded // Compute grad_norm (sum of squares) self.launch_grad_norm()?; @@ -1693,6 +1699,16 @@ impl GpuDqnTrainer { let (grad_norm_kernel, adam_update_kernel, f32_to_bf16_kernel, bf16_to_f32_kernel, saxpy_kernel, zero_kernel, regime_scale_kernel, shrink_perturb, _relu_mask_in_module, spectral_norm_kernel, clipped_saxpy_kernel, clip_grad_kernel) = compile_training_kernels(&stream, &config)?; + // Separate grad_norm instance for non-graph launches (clip_grad_buf_inplace). + // CUDA graph captures CUfunction handles — the same handle can't be launched + // both inside a replayed graph and outside on the same stream. + let grad_norm_standalone = { + let module = stream.context().load_cubin(DQN_UTILITY_CUBIN.to_vec()) + .map_err(|e| MLError::ModelError(format!("grad_norm_standalone cubin: {e}")))?; + module.load_function("dqn_grad_norm_kernel") + .map_err(|e| MLError::ModelError(format!("grad_norm_standalone load: {e}")))? + }; + // ── Compile EMA kernel (standalone — not captured in CUDA Graph) ── let ema_kernel = compile_ema_kernel(&stream)?; let relu_mask_standalone = compile_relu_mask_standalone(&stream)?; @@ -2006,6 +2022,7 @@ impl GpuDqnTrainer { stream, ptrs, grad_norm_kernel, + grad_norm_standalone, adam_update_kernel, ema_kernel, per_update_kernel, @@ -3554,6 +3571,7 @@ impl GpuDqnTrainer { fn launch_grad_norm(&self) -> Result<(), MLError> { let tp = self.total_params as i32; let blocks = ((self.total_params + 255) / 256) as u32; + tp, blocks, self.ptrs.grad_buf, self.ptrs.grad_norm_buf); let launch_cfg = LaunchConfig { grid_dim: (blocks, 1, 1), @@ -3568,7 +3586,7 @@ impl GpuDqnTrainer { // grad_buf has size = total_params; grad_norm_buf has size = 1. unsafe { self.stream - .launch_builder(&self.grad_norm_kernel) + .launch_builder(&self.grad_norm_standalone) .arg(&grad_ptr) .arg(&norm_ptr) .arg(&tp) diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index a5e316601..fba98d0c2 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -602,6 +602,7 @@ impl FusedTrainingCtx { let c51_frac = 1.0 - cql_frac - iqn_frac - ens_frac; let c51_budget = self.trainer.config().max_grad_norm * c51_frac; + self.trainer.clip_grad_buf_inplace(c51_budget) .map_err(|e| anyhow::anyhow!("C51 gradient budget clip: {e}"))?; }