feat(cuda): bypass check_gradients_finite when GPU training guard active

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

crates/ml-core/src/gradient_accumulation.rs

@@ -138,6 +138,23 @@ pub fn check_gradients_finite(grads: &GradStore, vars: &[Var]) -> Result<(), MLE
     crate::gradient_utils::check_gradients_finite(vars, grads)
 }
 
+/// Check gradients are finite -- skips GPU sync when training guard is active.
+///
+/// When `gpu_guard_active` is true, the GPU training guard kernel has already
+/// checked for NaN/Inf on the loss and grad_norm scalars. The per-parameter
+/// NaN check is redundant: if any gradient were NaN, the norm (from
+/// `clip_grad_norm`) would also be NaN, which the guard catches.
+pub fn check_gradients_finite_guarded(
+    grads: &GradStore,
+    vars: &[Var],
+    gpu_guard_active: bool,
+) -> Result<(), MLError> {
+    if gpu_guard_active {
+        return Ok(());
+    }
+    crate::gradient_utils::check_gradients_finite(vars, grads)
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

crates/ml/src/trainers/dqn/trainer.rs

@@ -4825,8 +4825,15 @@ impl DQNTrainer {
             )
             .map_err(|e| anyhow::anyhow!("Gradient scaling failed: {}", e))?;
 
-            crate::gradient_accumulation::check_gradients_finite(grads, &vars)
-                .map_err(|e| anyhow::anyhow!("Training halted: {}", e))?;
+            #[cfg(feature = "cuda")]
+            let guard_active = self.training_guard.is_some();
+            #[cfg(not(feature = "cuda"))]
+            let guard_active = false;
+            crate::gradient_accumulation::check_gradients_finite_guarded(
+                grads,
+                &vars,
+                guard_active,
+            ).map_err(|e| anyhow::anyhow!("Training halted: {}", e))?;
 
             agent
                 .apply_accumulated_gradients(grads)