ruda-optim 0.21.13

Ruda optimizer updates, gradient state, clipping and learning-rate schedules.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37

use super::GradientsParams;
use alloc::vec::Vec;
use ruda_model::tensor::{TensorPrimitive, backend::Backend};
use ruccl::{
    ReduceOperation,
    rank::communicator::RankCommunicator,
    tensor_device::{TensorDevice, TensorDeviceError},
};

impl GradientsParams {
    /// Synchronize gradients using an explicitly owned ruCCL rank communicator.
    ///
    /// Peers must have matching parameter IDs, gradient shapes and dtypes, and
    /// invoke this operation in the same order. Missing gradients are not
    /// replaced with zeros. Use the inner backend when training with autodiff.
    pub fn all_reduce_with<B: Backend>(
        mut self,
        communicator: &RankCommunicator<TensorDevice<B>>,
        operation: ReduceOperation,
    ) -> Result<Self, TensorDeviceError> {
        let mut ids = self.container.ids().into_iter().copied().collect::<Vec<_>>();
        ids.sort();
        for id in ids {
            let gradient = self.container.remove::<B>(&id).ok_or(
                TensorDeviceError::InvalidBuffer("gradient parameter is missing"),
            )?;
            let TensorPrimitive::Float(gradient) = gradient else {
                return Err(TensorDeviceError::InvalidOperation(
                    "quantized gradients cannot be all-reduced",
                ));
            };
            let gradient = communicator.all_reduce_float(gradient, operation)?;
            self.container.register::<B>(id, TensorPrimitive::Float(gradient));
        }
        Ok(self)
    }
}