diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 80ab9f567..4452861fa 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -2053,8 +2053,11 @@ impl GpuDqnTrainer { // Disable cudarc event tracking — graph replay leaves stale events // that cause INVALID_VALUE on post-graph kernel launches. let _evt_guard = EventTrackingGuard::new(self.stream.context()); - // Two-phase grad_norm — no memset needed - self.launch_grad_norm()?; + // Two-phase grad_norm using STANDALONE kernel handle (NOT the captured one). + // After graph_mega capture, the captured grad_norm_kernel CUfunction cannot be + // launched outside the graph context — CUDA corrupts kernel state silently. + // grad_norm_standalone is a separately loaded handle of the same kernel. + self.launch_grad_norm_standalone()?; self.launch_grad_norm_finalize()?; // Clip in-place @@ -4697,10 +4700,18 @@ impl GpuDqnTrainer { /// `alpha = 1.0`: pure C51 distributional loss (converged). /// Intermediate values blend: `grad = alpha*C51 + (1-alpha)*MSE`. /// - /// No CUDA Graph invalidation — the kernel sequence is always the same - /// (both MSE and C51 run every step); only the SAXPY scale factors change. + /// INVALIDATES graph_mega — alpha/scale_mse are passed as literal scalar + /// arguments to the SAXPY blend kernels inside the graph. CUDA Graph bakes + /// scalar args at capture time. Without invalidation, the blend is frozen + /// at the capture-time alpha forever. The graph is re-captured on next step. pub fn set_c51_alpha(&mut self, alpha: f32) { - self.c51_alpha = alpha.clamp(0.0, 1.0); + let new_alpha = alpha.clamp(0.0, 1.0); + if (new_alpha - self.c51_alpha).abs() > 1e-6 { + self.graph_forward = None; + self.graph_forward_ddqn = None; + // graph_mega invalidated via the caller (FusedTrainingCtx) + } + self.c51_alpha = new_alpha; } /// Blend MSE and C51 loss values with NaN-safe alpha guard. @@ -5640,6 +5651,36 @@ impl GpuDqnTrainer { Ok(()) } + /// Launch grad norm using the STANDALONE kernel handle (not the captured one). + /// Must be used outside CUDA Graph context (e.g. clip_grad_buf_inplace). + /// After graph_mega capture, the captured CUfunction handle cannot be launched + /// outside the graph — CUDA silently corrupts kernel state. + fn launch_grad_norm_standalone(&self) -> Result<(), MLError> { + let tp = self.total_params as i32; + let blocks = self.grad_norm_blocks as u32; + + let grad_ptr = self.ptrs.grad_buf; + let partials_ptr = self.grad_norm_partials.raw_ptr(); + + unsafe { + self.stream + .launch_builder(&self.grad_norm_standalone) + .arg(&grad_ptr) + .arg(&partials_ptr) + .arg(&tp) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| { + MLError::ModelError(format!("dqn_grad_norm_standalone launch: {e}")) + })?; + } + + Ok(()) + } + /// Launch the Adam update kernel with correct gradient clipping. /// /// Argument order matches `dqn_adam_update_kernel` in `dqn_training_kernel.cu`: diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 5e155fc8d..026246780 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -1626,9 +1626,13 @@ impl FusedTrainingCtx { /// Set C51 blend factor for gradual MSE→C51 ramp. /// /// `alpha = 0.0`: pure MSE (warmup). `alpha = 1.0`: pure C51 (converged). - /// No graph invalidation — kernel sequence is fixed (both always run). + /// Invalidates graph_mega — alpha is a scalar arg baked at capture time. pub(crate) fn set_c51_alpha(&mut self, alpha: f32) { + let prev = self.trainer.c51_alpha(); self.trainer.set_c51_alpha(alpha); + if (alpha - prev).abs() > 1e-6 { + self.graph_mega = None; // Force recapture with new alpha + } } /// Return a reference to the GPU-resident loss scalar (f32) for the training guard.