foxhunt/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(())
}

/// 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.0f32, 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_nan_gradients_detected() {
        let device = Device::Cpu;
        let var = Var::from_tensor(
            &candle_core::Tensor::new(&[0.0f32], &device).expect("failed to create tensor"),
        )
        .expect("failed to create var");
        let zero = candle_core::Tensor::new(&[0.0f32], &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
        );
    }
}