foxhunt/crates/ml/src/gradient_accumulation.rs

//! Gradient accumulation utilities for mini-batch training
//!
//! Provides functions to accumulate gradients across multiple micro-batches
//! and scale them before applying an optimizer step. This is essential for
//! training with effective batch sizes larger than what fits in GPU memory.
//!
//! # Usage
//!
//! ```no_run
//! use ml::gradient_accumulation::{accumulate_grads, scale_grads};
//! use candle_core::backprop::GradStore;
//! use candle_core::Var;
//!
//! let vars: Vec<Var> = vec![]; // Your model variables
//! let accumulation_steps = 4;
//! let mut accumulated: Option<GradStore> = None;
//!
//! for _step in 0..accumulation_steps {
//!     // let grads = loss.backward()?;  // compute micro-batch grads
//!     // accumulate_grads(&mut accumulated, grads, &vars)?;
//! }
//!
//! // Scale by 1/accumulation_steps before optimizer step
//! // if let Some(ref mut grads) = accumulated {
//! //     scale_grads(grads, &vars, 1.0 / accumulation_steps as f64)?;
//! // }
//! ```

use candle_core::backprop::GradStore;
use candle_core::Var;

use crate::MLError;

/// Accumulate gradients from a micro-batch into an accumulator.
///
/// On the first call (when `target` is `None`), the source gradients become
/// the accumulator. On subsequent calls, source gradients are added element-wise
/// to the existing accumulator.
///
/// # Arguments
///
/// * `target` - Mutable reference to the accumulator. `None` on first call.
/// * `source` - Gradients from the current micro-batch (consumed).
/// * `vars` - Slice of model variables whose gradients should be accumulated.
///
/// # Errors
///
/// Returns an error if tensor addition fails (e.g., shape mismatch).
pub fn accumulate_grads(
    target: &mut Option<GradStore>,
    source: GradStore,
    vars: &[Var],
) -> Result<(), MLError> {
    match target {
        None => {
            // First micro-batch: source becomes the accumulator
            *target = Some(source);
        }
        Some(ref mut acc) => {
            // Subsequent micro-batches: add element-wise
            for var in vars {
                if let Some(src_grad) = source.get(var) {
                    if let Some(existing) = acc.remove(var) {
                        let summed = existing.add(src_grad).map_err(|e| {
                            MLError::TrainingError(format!(
                                "Failed to accumulate gradients: {}",
                                e
                            ))
                        })?;
                        acc.insert(var, summed);
                    } else {
                        // Variable had no gradient in accumulator yet; clone source
                        let cloned = src_grad.clone();
                        acc.insert(var, cloned);
                    }
                }
            }
        }
    }
    Ok(())
}

/// Scale all gradients in a `GradStore` by a scalar factor.
///
/// This is typically used after accumulation to divide by the number of
/// accumulation steps, producing the mean gradient.
///
/// # Arguments
///
/// * `grads` - Mutable reference to the gradient store to scale in-place.
/// * `vars` - Slice of model variables whose gradients should be scaled.
/// * `scale` - The scalar multiplier (e.g., `1.0 / accumulation_steps as f64`).
///
/// # Errors
///
/// Returns an error if tensor multiplication fails.
pub fn scale_grads(grads: &mut GradStore, vars: &[Var], scale: f64) -> Result<(), MLError> {
    for var in vars {
        if let Some(grad) = grads.remove(var) {
            let scaled = (grad * scale).map_err(|e| {
                MLError::TrainingError(format!("Failed to scale gradient: {}", e))
            })?;
            grads.insert(var, scaled);
        }
    }
    Ok(())
}

/// Clip gradients by global L2 norm.
///
/// If the total L2 norm of all gradients exceeds `max_norm`, all gradients
/// are scaled down proportionally so the total norm equals `max_norm`.
/// This is equivalent to PyTorch's `torch.nn.utils.clip_grad_norm_`.
///
/// Returns the original (pre-clip) gradient norm for logging.
pub fn clip_grads(grads: &mut GradStore, vars: &[Var], max_norm: f64) -> Result<f64, MLError> {
    // Compute total L2 norm
    let mut total_norm_sq = 0.0_f64;
    for var in vars {
        if let Some(grad) = grads.get(var) {
            let norm_sq = grad
                .sqr()
                .map_err(|e| {
                    MLError::TrainingError(format!("Failed to square gradient for clipping: {}", e))
                })?
                .sum_all()
                .map_err(|e| {
                    MLError::TrainingError(format!("Failed to sum gradient for clipping: {}", e))
                })?
                .to_vec0::<f32>()
                .map_err(|e| {
                    MLError::TrainingError(format!("Failed to extract norm for clipping: {}", e))
                })? as f64;
            total_norm_sq += norm_sq;
        }
    }
    let total_norm = total_norm_sq.sqrt();

    if total_norm > max_norm && total_norm > 0.0 {
        let clip_coef = max_norm / total_norm;
        scale_grads(grads, vars, clip_coef)?;
    }

    Ok(total_norm)
}

/// Check if any gradient in the GradStore contains NaN or Inf.
/// Returns `Err` with the variable index if any gradient is non-finite.
pub fn check_gradients_finite(grads: &GradStore, vars: &[Var]) -> Result<(), MLError> {
    for (idx, var) in vars.iter().enumerate() {
        if let Some(grad) = grads.get(var) {
            let flat = grad.flatten_all().map_err(|e| {
                MLError::TrainingError(format!("Failed to flatten gradient {}: {}", idx, e))
            })?;
            let values = flat.to_vec1::<f32>().map_err(|e| {
                MLError::TrainingError(format!("Failed to read gradient {}: {}", idx, e))
            })?;
            for val in &values {
                if !val.is_finite() {
                    return Err(MLError::TrainingError(format!(
                        "NaN/Inf gradient detected in parameter {} (shape: {:?}). \
                         Halting training to prevent model corruption.",
                        idx,
                        grad.shape()
                    )));
                }
            }
        }
    }
    Ok(())
}

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

    #[test]
    fn test_finite_gradients_pass() {
        let device = Device::Cpu;
        let var = Var::from_tensor(
            &candle_core::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(&grads, &[var]);
        assert!(result.is_ok());
    }

    #[test]
    fn test_clip_grads_reduces_norm() {
        let device = Device::Cpu;
        // Create a var with large values to produce large gradients
        let var = Var::from_tensor(
            &candle_core::Tensor::new(&[10.0_f32, 20.0, 30.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 mut grads = loss.backward().expect("backward failed");

        // Gradients are [20, 40, 60], norm = sqrt(20^2 + 40^2 + 60^2) = sqrt(5600) ≈ 74.8
        let max_norm = 1.0;
        let orig_norm = clip_grads(&mut grads, &[var.clone()], max_norm).expect("clip failed");
        assert!(
            orig_norm > max_norm,
            "Original norm {orig_norm} should exceed max_norm {max_norm}"
        );

        // After clipping, verify the clipped gradient norm is ≈ max_norm
        let clipped_grad = grads.get(&var).expect("gradient should exist");
        let clipped_norm_sq = clipped_grad
            .sqr()
            .expect("sqr failed")
            .sum_all()
            .expect("sum failed")
            .to_vec0::<f32>()
            .expect("extract failed") as f64;
        let clipped_norm = clipped_norm_sq.sqrt();
        assert!(
            (clipped_norm - max_norm).abs() < 0.01,
            "Clipped norm {clipped_norm} should be close to max_norm {max_norm}"
        );
    }

    #[test]
    fn test_clip_grads_noop_when_below_max() {
        let device = Device::Cpu;
        let var = Var::from_tensor(
            &candle_core::Tensor::new(&[0.1_f32, 0.2], &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 mut grads = loss.backward().expect("backward failed");

        // Gradients are [0.2, 0.4], norm ≈ 0.447 — below max_norm
        let max_norm = 5.0;
        let orig_norm = clip_grads(&mut grads, &[var.clone()], max_norm).expect("clip failed");
        assert!(
            orig_norm < max_norm,
            "Norm {orig_norm} should be below max {max_norm}"
        );

        // Gradient should be unchanged
        let grad = grads.get(&var).expect("gradient should exist");
        let values = grad.to_vec1::<f32>().expect("extract failed");
        assert!((values[0] - 0.2).abs() < 0.001, "Gradient should be unchanged");
        assert!((values[1] - 0.4).abs() < 0.001, "Gradient should be unchanged");
    }

    #[test]
    fn test_nan_gradients_detected() {
        let device = Device::Cpu;
        let var = Var::from_tensor(
            &candle_core::Tensor::new(&[0.0_f32], &device).expect("failed to create tensor"),
        )
        .expect("failed to create var");
        let zero = candle_core::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(&grads, &[var]);
        assert!(result.is_err());
        let err_msg = format!("{}", result.expect_err("expected error"));
        assert!(
            err_msg.contains("NaN") || err_msg.contains("Inf"),
            "Expected NaN/Inf mention in: {}",
            err_msg
        );
    }
}