diff --git a/crates/ml-core/src/gradient_accumulation.rs b/crates/ml-core/src/gradient_accumulation.rs index 72324c1f2..27adc3f07 100644 --- a/crates/ml-core/src/gradient_accumulation.rs +++ b/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::*; diff --git a/crates/ml/src/trainers/dqn/trainer.rs b/crates/ml/src/trainers/dqn/trainer.rs index 56ae3f50f..a9f456be8 100644 --- a/crates/ml/src/trainers/dqn/trainer.rs +++ b/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)