foxhunt/crates/ml-core/src/gradient_utils.rs

//! Gradient utilities for Candle framework
//!
//! Provides gradient clipping, NaN/Inf detection, and other gradient-related
//! operations that are missing from candle_nn::optim.
//!
//! **GPU hot-path rule**: ZERO `to_scalar()` / `to_vec*()` calls in any
//! function on this module's critical path.  All norm/sum computations
//! stay on GPU-resident tensors. Callers that need a CPU scalar must
//! call `Tensor::to_scalar()` themselves — typically AFTER `optimizer.step()`
//! so the readback piggybacks on an already-flushed pipeline.

use candle_core::{backprop::GradStore, DType, Error, Tensor, TensorId, Var};

use crate::MLError;

/// Throttle counter for the TensorId fallback warning.
/// Logs on first occurrence and then every 1000 calls to avoid spam.
static FALLBACK_LOG_COUNT: std::sync::atomic::AtomicUsize =
    std::sync::atomic::AtomicUsize::new(0);

/// Clip gradients by global L2 norm — **fully GPU-resident, zero CPU sync**.
///
/// Computes the L2 norm of all gradients on GPU, then unconditionally
/// multiplies every gradient by `min(max_norm / (norm + eps), 1.0)`.
/// When the norm is already below `max_norm` the scale factor is clamped
/// to 1.0, making the multiply a no-op in terms of gradient values.
///
/// **Var identity resilience**: If the provided `vars` don't match any
/// GradStore entries (Var/TensorId mismatch), the function falls back
/// to iterating GradStore entries by `TensorId` directly — computing
/// the norm AND applying clipping via `insert_id`. This guarantees
/// gradient clipping works regardless of Var object identity.
///
/// **Why unconditional multiply?**  A conditional `if norm > max_norm`
/// requires reading the norm to CPU (`to_scalar()`), which inserts a
/// `cuStreamSynchronize` barrier between backward() and optimizer.step().
/// The unconditional GPU multiply costs < 1µs and keeps the pipeline
/// fully saturated.
///
/// # Returns
/// A `[1]`-shaped F32 tensor containing the pre-clip gradient norm,
/// resident on the **same device** as the gradients.  The caller can
/// read it with `to_scalar::<f32>()` at a convenient sync boundary
/// (e.g. after `optimizer.step()` + `loss.to_scalar()`).
pub fn clip_grad_norm(
    vars: &[Var],
    grads: &mut GradStore,
    max_norm: f64,
    device: &candle_core::Device,
) -> Result<Tensor, Error> {

    // Accumulate squared norms on a GPU-resident scalar tensor.
    // ZERO to_scalar() calls — all arithmetic stays on device.
    //
    // Strategy: try var-based matching first (fast path). If zero vars match,
    // fall back to iterating all GradStore entries by TensorId (always correct).
    let mut norm_sq_acc = Tensor::zeros(&[1], DType::F32, device)?;
    let mut matched = 0_usize;
    for var in vars.iter() {
        if let Some(grad) = grads.get(var) {
            let partial = grad.sqr()?.sum_all()?.to_dtype(DType::F32)?.reshape(&[1])?;
            norm_sq_acc = (norm_sq_acc + partial)?;
            matched += 1;
        }
    }

    // Fallback: compute norm from all GradStore entries directly.
    // Covers Var identity mismatch (e.g. optimizer holds stale clones).
    let use_id_path = matched == 0 && !vars.is_empty();
    if use_id_path {
        // Throttled warning: first call + every 1000th to avoid log spam
        let count = FALLBACK_LOG_COUNT.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        if count == 0 || count % 1000 == 0 {
            tracing::warn!(
                "clip_grad_norm: 0/{} vars matched GradStore — using TensorId fallback (occurrence #{})",
                vars.len(),
                count + 1,
            );
        }

        for id in grads.get_ids() {
            if let Some(grad) = grads.get_id(*id) {
                let partial = grad.sqr()?.sum_all()?.to_dtype(DType::F32)?.reshape(&[1])?;
                norm_sq_acc = (norm_sq_acc + partial)?;
            }
        }
    }

    // GPU-resident norm (never leaves the device)
    let norm = norm_sq_acc.sqrt()?;

    // scale = min(max_norm / (norm + eps), 1.0)  — entirely on GPU
    let max_norm_t = Tensor::new(&[max_norm as f32], device)?;
    let eps_t = Tensor::new(&[1e-6_f32], device)?;
    let scale = (&max_norm_t / (&norm + &eps_t)?)?
        .clamp(0.0_f64, 1.0)?;

    // Unconditionally multiply every gradient by `scale`.
    // When norm ≤ max_norm, scale ≈ 1.0 so values are unchanged.
    // Cast scale to each gradient's dtype to handle mixed-precision (BF16) models.
    if use_id_path {
        // TensorId path: clip via insert_id — bypasses Var identity entirely.
        // Collect IDs first to avoid borrowing `grads` mutably while iterating.
        let ids: Vec<TensorId> = grads.get_ids().copied().collect();
        for id in ids {
            if let Some(grad) = grads.get_id(id) {
                let scale_cast = scale.to_dtype(grad.dtype())?;
                let scaled_grad = grad.broadcast_mul(&scale_cast)?;
                grads.insert_id(id, scaled_grad);
            }
        }
    } else {
        for var in vars.iter() {
            if let Some(grad) = grads.remove(var) {
                let scale_cast = scale.to_dtype(grad.dtype())?;
                let scaled_grad = grad.broadcast_mul(&scale_cast)?;
                grads.insert(var, scaled_grad);
            }
        }
    }

    Ok(norm)
}

/// Check if any gradient in the GradStore contains NaN or Inf values.
///
/// Accumulates a validity flag on GPU — single `to_scalar()` at the end.
/// Previous implementation called `to_scalar()` per parameter tensor.
///
/// Returns `Ok(())` if all gradients are finite, or an error if any
/// gradient contains NaN or Inf values.
pub fn check_gradients_finite(vars: &[Var], grads: &GradStore) -> Result<(), MLError> {
    let device = vars.iter()
        .find_map(|v| grads.get(v).map(|g| g.device().clone()))
        .unwrap_or(candle_core::Device::Cpu);

    // Accumulate sum-of-sums on GPU. If ANY element is NaN, the total is NaN.
    let mut total_sum = Tensor::zeros(&[], DType::F32, &device)
        .map_err(|e| MLError::TrainingError(format!("Failed to create accumulator: {}", e)))?;

    for var in vars.iter() {
        if let Some(grad) = grads.get(var) {
            let partial = grad.sum_all()
                .and_then(|t| t.to_dtype(DType::F32))
                .map_err(|e| {
                    MLError::TrainingError(format!("Failed to sum gradient: {}", e))
                })?;
            total_sum = (total_sum + partial).map_err(|e| {
                MLError::TrainingError(format!("Failed to accumulate gradient sum: {}", e))
            })?;
        }
    }

    // Single GPU→CPU sync
    let sum_val = total_sum.to_scalar::<f32>().map_err(|e| {
        MLError::TrainingError(format!("Failed to read gradient sum: {}", e))
    })?;

    if !sum_val.is_finite() {
        return Err(MLError::TrainingError(
            "NaN/Inf gradient detected in model parameters. \
             Training is numerically unstable -- halting to prevent model corruption. \
             Consider reducing learning rate or checking input data for anomalies."
                .to_owned(),
        ));
    }

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use candle_core::{Device, Tensor};

    #[test]
    fn test_finite_gradients_pass() {
        let device = Device::Cpu;
        let var = Var::from_tensor(
            &Tensor::new(&[1.0_f32, 2.0, 3.0], &device).expect("failed to create tensor"),
        )
        .expect("failed to create var");

        let loss = var
            .mul(&var)
            .expect("mul failed")
            .sum_all()
            .expect("sum failed");
        let grads = loss.backward().expect("backward failed");

        let result = check_gradients_finite(&[var], &grads);
        assert!(result.is_ok(), "Finite gradients should pass check");
    }

    #[test]
    fn test_nan_gradient_detected() {
        let device = Device::Cpu;
        let var = Var::from_tensor(
            &Tensor::new(&[0.0_f32], &device).expect("failed to create tensor"),
        )
        .expect("failed to create var");

        // 0/0 produces NaN
        let zero = Tensor::new(&[0.0_f32], &device).expect("failed to create zero");
        let nan_result = var.div(&zero).expect("div failed");
        let loss = nan_result.sum_all().expect("sum failed");
        let grads = loss.backward().expect("backward failed");

        let result = check_gradients_finite(&[var], &grads);
        assert!(result.is_err(), "NaN/Inf gradients should be detected");
        let err_msg = format!("{}", result.expect_err("expected error"));
        assert!(
            err_msg.contains("NaN/Inf gradient detected"),
            "Error should mention NaN/Inf: {}",
            err_msg
        );
    }

    #[test]
    fn test_empty_vars_passes() {
        let device = Device::Cpu;
        // Create an empty GradStore by computing backward on a constant
        let loss = Tensor::new(1.0_f32, &device).expect("failed to create tensor");
        let grads = loss.backward().expect("backward failed");

        let vars: Vec<Var> = vec![];
        let result = check_gradients_finite(&vars, &grads);
        assert!(result.is_ok(), "Empty vars should pass");
    }

    #[test]
    fn test_clip_grad_norm_below_threshold() {
        let device = Device::Cpu;
        let var = Var::from_tensor(
            &Tensor::new(&[1.0_f32, 2.0, 3.0], &device).expect("create"),
        )
        .expect("var");
        let loss = var.mul(&var).expect("mul").sum_all().expect("sum");
        let mut grads = loss.backward().expect("backward");

        // grad = [2, 4, 6], norm = sqrt(4+16+36) = sqrt(56) ≈ 7.48
        let norm_tensor = clip_grad_norm(&[var.clone()], &mut grads, 100.0, &device).expect("clip");
        let actual = norm_tensor.to_vec1::<f32>().expect("read")[0] as f64;
        assert!((actual - 7.483_f64).abs() < 0.01);
        // Gradients should be unchanged (norm < max_norm → scale ≈ 1.0)
        let grad = grads.get(&var).expect("grad exists");
        let g_vec = grad.to_vec1::<f32>().expect("read");
        assert!((g_vec[0] - 2.0).abs() < 0.02, "grad[0] should be ~2.0, got {}", g_vec[0]);
    }

    #[test]
    fn test_clip_grad_norm_above_threshold() {
        let device = Device::Cpu;
        let var = Var::from_tensor(
            &Tensor::new(&[1.0_f32, 2.0, 3.0], &device).expect("create"),
        )
        .expect("var");
        let loss = var.mul(&var).expect("mul").sum_all().expect("sum");
        let mut grads = loss.backward().expect("backward");

        // grad = [2, 4, 6], norm ≈ 7.48, clip to 1.0
        let norm_tensor = clip_grad_norm(&[var.clone()], &mut grads, 1.0, &device).expect("clip");
        let actual = norm_tensor.to_vec1::<f32>().expect("read")[0] as f64;
        assert!(actual > 1.0);
        // Gradients should be scaled down: new_norm ≈ 1.0
        let grad = grads.get(&var).expect("grad exists");
        let g_vec = grad.to_vec1::<f32>().expect("read");
        let new_norm = (g_vec[0] * g_vec[0] + g_vec[1] * g_vec[1] + g_vec[2] * g_vec[2]).sqrt();
        assert!((new_norm - 1.0).abs() < 0.05, "Clipped norm should be ~1.0, got {}", new_norm);
    }

    /// Test that the TensorId fallback fires when vars don't match the GradStore.
    /// Simulates the Var identity mismatch: create loss from var_a, but pass var_b
    /// (a different Var with a different TensorId) to clip_grad_norm.
    #[test]
    fn test_clip_grad_norm_var_mismatch_fallback() {
        let device = Device::Cpu;
        // var_a is used in the forward pass
        let var_a = Var::from_tensor(
            &Tensor::new(&[1.0_f32, 2.0, 3.0], &device).expect("create"),
        )
        .expect("var_a");
        let loss = var_a.mul(&var_a).expect("mul").sum_all().expect("sum");
        let mut grads = loss.backward().expect("backward");

        // var_b is a DIFFERENT Var (different TensorId) — simulates the mismatch
        let var_b = Var::from_tensor(
            &Tensor::new(&[0.0_f32, 0.0, 0.0], &device).expect("create"),
        )
        .expect("var_b");

        // clip_grad_norm with mismatched vars should still compute correct norm
        // via the TensorId fallback
        let norm_tensor = clip_grad_norm(&[var_b], &mut grads, 100.0, &device).expect("clip");
        let actual = norm_tensor.to_vec1::<f32>().expect("read")[0] as f64;
        // grad of var_a = [2, 4, 6], norm = sqrt(56) ≈ 7.48
        assert!(actual > 7.0, "Fallback norm should be ~7.48, got {}", actual);
    }

    /// Test that the TensorId fallback ACTUALLY CLIPS gradients (not just computes norm).
    /// This was the critical bug: fallback computed correct norm but left gradients unclipped.
    #[test]
    fn test_clip_grad_norm_var_mismatch_actually_clips() {
        let device = Device::Cpu;
        let var_a = Var::from_tensor(
            &Tensor::new(&[1.0_f32, 2.0, 3.0], &device).expect("create"),
        )
        .expect("var_a");
        let loss = var_a.mul(&var_a).expect("mul").sum_all().expect("sum");
        let mut grads = loss.backward().expect("backward");

        // Mismatched var — forces TensorId fallback
        let var_b = Var::from_tensor(
            &Tensor::new(&[0.0_f32, 0.0, 0.0], &device).expect("create"),
        )
        .expect("var_b");

        // Clip to max_norm=1.0 — grad norm is ~7.48, must be scaled down
        let norm_tensor = clip_grad_norm(&[var_b], &mut grads, 1.0, &device).expect("clip");
        let pre_clip_norm = norm_tensor.to_vec1::<f32>().expect("read")[0] as f64;
        assert!(pre_clip_norm > 1.0, "Pre-clip norm should be >1.0, got {}", pre_clip_norm);

        // Verify the gradient was ACTUALLY clipped by reading it via TensorId
        let var_a_id = var_a.as_tensor().id();
        let clipped_grad = grads.get_id(var_a_id).expect("clipped grad should exist");
        let g_vec = clipped_grad.to_vec1::<f32>().expect("read");
        let clipped_norm = (g_vec[0] * g_vec[0] + g_vec[1] * g_vec[1] + g_vec[2] * g_vec[2]).sqrt();
        assert!(
            (clipped_norm - 1.0).abs() < 0.05,
            "Clipped norm should be ~1.0, got {} (grads: {:?})",
            clipped_norm, g_vec
        );
    }
}