108 lines
3.7 KiB
Rust
108 lines
3.7 KiB
Rust
//! 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(())
|
|
}
|