From 03ad33fda4aa4814e7d9d8a2d627160df5a520d0 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sat, 28 Mar 2026 15:50:38 +0100 Subject: [PATCH] fix(bf16): EventTrackingGuard for post-graph kernel launches MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Root cause: cudarc event tracking left stale state after CUDA graph replay, causing INVALID_VALUE on subsequent kernel launches. Fix: add EventTrackingGuard to clip_grad_buf_inplace (same pattern as read_grad_norm_sync, apply_spectral_norm, etc.). Smoke test now runs 3 epochs / 600 steps successfully. Q-values diverge (bf16 precision) — needs hyperparameter tuning for bf16 (lower lr, smaller Huber delta, tighter gradient clip). Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 1886e9042..feb3a8e4f 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -1521,7 +1521,9 @@ 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) + // 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()); // Zero the grad_norm accumulator self.stream .memset_zeros(&mut self.grad_norm_buf) @@ -3571,7 +3573,6 @@ 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),