From b7ea7c07600f599e8ce14f912dc56866610caec9 Mon Sep 17 00:00:00 2001
From: jgrusewski <jgrusewski@localhost>
Date: Sun, 5 Apr 2026 15:33:36 +0200
Subject: [PATCH] chore: remove runtime NaN tracing stream syncs
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Removes run_nan_tracing (3 stream syncs per step on first 20 steps)
and unused grad_buf_ptr/params_buf_ptr accessors. Root causes are
fixed — runtime detection no longer needed. isfinite guards in CUDA
kernels remain as silent defense-in-depth.

17/17 smoke tests pass.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs   | 14 ------
 crates/ml/src/trainers/dqn/fused_training.rs  | 45 -------------------
 2 files changed, 59 deletions(-)

diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
index 7af63269f..e2187a142 100644
--- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
+++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
@@ -1000,15 +1000,6 @@ impl GpuDqnTrainer {
 
     /// Raw device pointer to `grad_buf` for auxiliary gradient injection between graph phases.
     ///
-    /// Between `graph_forward` replay (which populates `grad_buf` with C51 gradients)
-    /// and `graph_adam` replay (which reads `grad_buf` for Adam), external code can ADD
-    /// auxiliary gradients (IQN trunk, ensemble heads) to `grad_buf` via SAXPY.
-    ///
-    /// Caller must use `EventTrackingGuard` and operate on the same stream.
-    pub fn grad_buf_ptr(&self) -> u64 {
-        self.grad_buf.raw_ptr()
-    }
-
     /// Mutable reference to grad_buf for test injection.
     #[cfg(test)]
     pub fn grad_buf_mut(&mut self) -> &mut CudaSlice<f32> {
@@ -1021,11 +1012,6 @@ impl GpuDqnTrainer {
         &mut self.cql_grad_scratch
     }
 
-    /// Raw pointer to f32 master params buffer (for NaN diagnostics).
-    pub fn params_buf_ptr(&self) -> u64 {
-        self.params_buf.raw_ptr()
-    }
-
     /// Total number of trainable parameters.
     pub fn total_params(&self) -> usize {
         self.total_params
diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs
index 835093e69..672d53ae1 100644
--- a/crates/ml/src/trainers/dqn/fused_training.rs
+++ b/crates/ml/src/trainers/dqn/fused_training.rs
@@ -655,9 +655,6 @@ impl FusedTrainingCtx {
         let gpu_batch = batch.gpu_batch.as_ref()
             .ok_or_else(|| anyhow::anyhow!("Fused training requires gpu_batch (GPU PER)"))?;
 
-        // ── NaN detection: GPU-side checks on first 20 steps ────────────
-        let nan_diag = self.steps_since_varmap_sync < 20;
-
         // ── Step 1: Spectral normalization BEFORE forward pass ─────────
         self.trainer.apply_spectral_norm(&mut self.online_dueling, &mut self.online_branching)
             .map_err(|e| anyhow::anyhow!("Spectral norm (pre-forward): {e}"))?;
@@ -695,11 +692,6 @@ impl FusedTrainingCtx {
             }
         }
 
-        // ── NaN check: grad_buf after forward (before aux ops touch it) ──
-        if nan_diag {
-            self.run_nan_tracing("post_forward")?;
-        }
-
         // ── Step 3: Auxiliary ops (graph_aux captured or replayed) ────────
         if self.graph_aux.is_some() {
             self.pre_replay_state_update(agent);
@@ -709,11 +701,6 @@ impl FusedTrainingCtx {
             self.submit_aux_ops(agent, gpu_batch)?;
         }
 
-        // ── NaN check: grad_buf after aux ops (before conditional ops) ──
-        if nan_diag {
-            self.run_nan_tracing("post_aux")?;
-        }
-
         // ── Step 4: Conditional ops (outside graph) ──────────────────────
 
         // Ensemble diversity (variable topology, complex forward passes).
@@ -750,11 +737,6 @@ impl FusedTrainingCtx {
         let fused_result = self.trainer.replay_adam_and_readback()
             .map_err(|e| anyhow::anyhow!("graph_adam replay: {e}"))?;
 
-        // ── NaN check: after Adam update ────────────────────────────────
-        if nan_diag {
-            self.run_nan_tracing("post_adam")?;
-        }
-
         // ── Step 6: PER priority update (outside graph) ──────────────────
         if let (Some(priorities_tensor), Some((alpha, epsilon))) =
             (agent.priorities_tensor()?, agent.per_alpha_epsilon())
@@ -937,33 +919,6 @@ impl FusedTrainingCtx {
         Ok(())
     }
 
-    /// GPU-side NaN tracing across the full training step.
-    ///
-    /// Checks grad_buf before Adam (to catch aux pipeline NaN) and f32 params
-    /// after Adam (to catch optimizer NaN). One stream sync per call (~5µs).
-    fn run_nan_tracing(&mut self, phase: &str) -> Result<()> {
-        let step = self.steps_since_varmap_sync;
-        let tp = self.trainer.total_params();
-
-        self.trainer.reset_nan_flags()?;
-        self.trainer.check_nan_f32(self.trainer.params_buf_ptr(), tp, 0)?;
-        self.trainer.check_nan_f32(self.trainer.grad_buf_ptr(), tp, 1)?;
-        let flags = self.trainer.read_nan_flags()?;
-
-        let params_nan = flags[0] != 0;
-        let grad_nan = flags[1] != 0;
-        if params_nan || grad_nan {
-            tracing::warn!(
-                step,
-                phase,
-                f32_params_nan = params_nan,
-                grad_buf_nan = grad_nan,
-                "NaN detected in training pipeline"
-            );
-        }
-        Ok(())
-    }
-
     /// Capture the auxiliary ops into a CUDA graph for zero-overhead replay.
     ///
     /// Called once after the first step succeeds (all sub-trainers initialized).