diff --git a/ml/src/lib.rs b/ml/src/lib.rs index a16562143..85dbbd3a5 100644 --- a/ml/src/lib.rs +++ b/ml/src/lib.rs @@ -49,6 +49,7 @@ )] // Import common types properly - NO ALIASES THAT CONFLICT! +use candle_core::backprop::GradStore; use candle_core::Tensor; use candle_core::Var; use candle_nn::Optimizer; // For Adam optimizer support @@ -234,6 +235,64 @@ impl Adam { Ok(clipped_grad_norm) } + /// Perform backward pass with gradient clipping but WITHOUT an optimizer step. + /// + /// This is useful for gradient accumulation workflows where you want to + /// accumulate clipped gradients across multiple micro-batches before + /// applying a single optimizer step. + /// + /// # Arguments + /// + /// * `loss` - The loss tensor to compute gradients from + /// * `max_norm` - Maximum allowed gradient norm (gradients will be clipped to this value) + /// + /// # Returns + /// + /// Returns `Ok((grads, clipped_norm))` on success, where `grads` is the clipped + /// `GradStore` and `clipped_norm` is the gradient norm after clipping. + /// + /// # Errors + /// + /// This function will return an error if: + /// - The backward pass fails to compute gradients + /// - Gradient clipping fails + pub fn backward_and_clip( + &self, + loss: &Tensor, + max_norm: f64, + ) -> Result<(GradStore, f64), MLError> { + // Compute gradients via backward pass + let mut grads = loss + .backward() + .map_err(|e| MLError::TrainingError(format!("Backward pass failed: {}", e)))?; + + // Apply gradient clipping in-place + let (_actual_norm, clipped_norm) = + crate::gradient_utils::clip_grad_norm(&self.vars, &mut grads, max_norm) + .map_err(|e| MLError::TrainingError(format!("Gradient clipping failed: {}", e)))?; + + Ok((grads, clipped_norm)) + } + + /// Apply pre-computed gradients to the optimizer. + /// + /// This is the companion to `backward_and_clip` for gradient accumulation + /// workflows. After accumulating and scaling gradients, call this method + /// to perform the optimizer step. + /// + /// # Arguments + /// + /// * `grads` - Pre-computed (and possibly accumulated/scaled) gradients + /// + /// # Errors + /// + /// This function will return an error if the optimizer step fails + pub fn apply_grads(&mut self, grads: &GradStore) -> Result<(), MLError> { + Optimizer::step(&mut self.optimizer, grads) + .map_err(|e| MLError::TrainingError(format!("Optimizer step failed: {}", e)))?; + Ok(()) + } + /// Compute the L2 norm of all gradients fn compute_gradient_norm( &self,