mlx_rs/optimizers/
mod.rs

1//! Trait and implementations for optimizers.
2
3#![deny(missing_docs)]
4
5use std::{
6    borrow::{Borrow, Cow},
7    collections::HashMap,
8    path::Path,
9    rc::Rc,
10};
11
12use crate::{
13    array,
14    error::{IoError, StateProjectionError, UnflattenError},
15    module::{FlattenedModuleParam, ModuleParameters},
16    utils::{StateProjection, Updatable},
17    Array,
18};
19
20mod adadelta;
21mod adafactor;
22mod adagrad;
23mod adam;
24mod adamax;
25mod adamw;
26mod lion;
27mod rmsprop;
28mod sgd;
29
30pub use adadelta::*;
31pub use adafactor::*;
32pub use adagrad::*;
33pub use adam::*;
34pub use adamax::*;
35pub use adamw::*;
36use itertools::Itertools;
37pub use lion::*;
38pub use rmsprop::*;
39pub use sgd::*;
40
41// Unfortunate workaround to implement Updatable for mutable references of
42// optimizers This is needed because of the orphan rule and lack of negative
43// trait bound, otherwise we would need to implement Updatable for every
44// `Module`
45macro_rules! impl_updatable_for_mut_optimizer {
46    ($optimizer:ty) => {
47        impl Updatable for &'_ mut $optimizer {
48            fn state_projection(&mut self) -> Result<StateProjection<'_>, StateProjectionError> {
49                <$optimizer as Updatable>::state_projection(&mut **self)
50            }
51        }
52    };
53}
54use impl_updatable_for_mut_optimizer;
55
56macro_rules! optimizer_updatable_state_methods {
57    () => {
58        fn state_projection(&mut self) -> Result<StateProjection<'_>, StateProjectionError> {
59            self.state_mut().state_projection()
60        }
61    };
62}
63use optimizer_updatable_state_methods;
64
65/// Type alias for common optimizer state.
66pub type State<T = Array> = HashMap<Rc<str>, T>;
67
68/// Trait for optimizer states.
69pub trait OptimizerState: Sized {
70    /// Error type for unflatten.
71    type UnflattenError: std::error::Error + Into<IoError>;
72
73    /// Declare all required and optional optimizer-state slots.
74    fn state_projection(&mut self) -> Result<StateProjection<'_>, StateProjectionError>;
75
76    /// Flatten present state entries in stable key order.
77    fn flatten(&mut self) -> Result<Vec<(Rc<str>, &Array)>, StateProjectionError> {
78        Ok(self.state_projection()?.into_entries().collect())
79    }
80
81    /// Flatten mutable present state entries in stable key order.
82    fn flatten_mut(&mut self) -> Result<Vec<(Rc<str>, &mut Array)>, StateProjectionError> {
83        Ok(self.state_projection()?.into_entries_mut().collect())
84    }
85
86    /// Unflatten an iterator of key-value pairs into the optimizer state.
87    fn unflatten<I, K>(input: I) -> Result<Self, Self::UnflattenError>
88    where
89        I: IntoIterator<Item = (K, Array)>,
90        K: Ord + AsRef<str> + Into<Rc<str>>;
91
92    /// Save the optimizer state to a safetensors file.
93    fn save_safetensors(&mut self, path: impl AsRef<Path>) -> Result<(), IoError> {
94        let state = self.flatten().map_err(IoError::StateProjection)?;
95        Array::save_safetensors(state, None, path)
96    }
97
98    /// Load the optimizer state from a safetensors file.
99    fn load_safetensors(&mut self, path: impl AsRef<Path>) -> Result<(), IoError> {
100        let loaded = Array::load_safetensors(path)?;
101        let unflattened = Self::unflatten(loaded).map_err(Into::into)?;
102
103        *self = unflattened;
104
105        Ok(())
106    }
107}
108
109impl OptimizerState for State {
110    type UnflattenError = std::convert::Infallible;
111
112    fn state_projection(&mut self) -> Result<StateProjection<'_>, StateProjectionError> {
113        let mut projection = StateProjection::new();
114        for (key, value) in self {
115            projection.required(key.clone(), value)?;
116        }
117        Ok(projection)
118    }
119
120    fn unflatten<I, K>(input: I) -> Result<Self, Self::UnflattenError>
121    where
122        Self: Sized,
123        I: IntoIterator<Item = (K, Array)>,
124        K: Ord + AsRef<str> + Into<Rc<str>>,
125    {
126        Ok(input.into_iter().map(|(k, v)| (k.into(), v)).collect())
127    }
128}
129
130impl OptimizerState for State<(Array, Array)> {
131    type UnflattenError = UnflattenError;
132
133    fn state_projection(&mut self) -> Result<StateProjection<'_>, StateProjectionError> {
134        let mut projection = StateProjection::new();
135        for (key, (first, second)) in self {
136            projection.required(format!("{key}.0"), first)?;
137            projection.required(format!("{key}.1"), second)?;
138        }
139        Ok(projection)
140    }
141
142    fn unflatten<I, K>(input: I) -> Result<Self, Self::UnflattenError>
143    where
144        Self: Sized,
145        I: IntoIterator<Item = (K, Array)>,
146        K: Ord + AsRef<str> + Into<Rc<str>>,
147    {
148        let mut state = State::new();
149        let iter = input
150            .into_iter()
151            .sorted_by(|a, b| a.0.as_ref().cmp(b.0.as_ref()))
152            .chunks(2);
153
154        for mut chunk in &iter {
155            let first = chunk.next().ok_or(UnflattenError::ExpectingNextPair)?;
156            let second = chunk.next().ok_or(UnflattenError::ExpectingNextPair)?;
157
158            // Check if the keys match up to the last dot and the suffix is 0 and 1 (should be already sorted)
159            let first_key = first.0.as_ref();
160            let second_key = second.0.as_ref();
161            if !first_key.ends_with(".0") || !second_key.ends_with(".1") {
162                return Err(UnflattenError::InvalidKey);
163            }
164            if first_key[..first_key.len() - 2] != second_key[..second_key.len() - 2] {
165                return Err(UnflattenError::InvalidKey);
166            }
167
168            let key = &first_key[..first_key.len() - 2];
169            let key: Rc<str> = Rc::from(key);
170            state.insert(key, (first.1, second.1));
171        }
172        Ok(state)
173    }
174}
175
176/// Trait for optimizers.
177pub trait Optimizer: Updatable {
178    /// State of the optimizer.
179    type State: OptimizerState;
180
181    /// Get the state of the optimizer.
182    fn state(&self) -> &Self::State;
183
184    /// Get the mutable state of the optimizer.
185    fn state_mut(&mut self) -> &mut Self::State;
186
187    /// Update a single parameter with the given gradient.
188    ///
189    /// The implementation should look up the state for the parameter using the key and update the
190    /// state and the parameter accordingly. The key is provided instead of the state because it
191    /// would otherwise create a mutable borrow conflict with the rest of the optimizer fields.
192    fn update_single(
193        &mut self,
194        key: &Rc<str>,
195        gradient: &Array,
196        parameter: &mut Array,
197    ) -> crate::error::Result<()>;
198
199    /// Apply the gradients to the parameters of the model and update the model with the new
200    /// parameters.
201    fn update<M>(
202        &mut self,
203        model: &mut M,
204        gradients: impl Borrow<FlattenedModuleParam>,
205    ) -> crate::error::Result<()>
206    where
207        M: ModuleParameters,
208    {
209        let mut parameters = model.parameters_mut().flatten();
210
211        for (key, gradient) in gradients.borrow().iter() {
212            if let Some(parameter) = parameters.get_mut(key) {
213                self.update_single(key, gradient, parameter)?;
214            }
215        }
216
217        Ok(())
218    }
219}
220
221/// Type alias for clipped gradients that is returned by `clip_grad_norm`.
222pub type MaybeClippedGrads<'a> = HashMap<Rc<str>, Cow<'a, Array>>;
223
224/// Clips the global norm of the gradients
225///
226/// This function ensures that the global norm of the gradients does not exceed
227/// `max_norm`. It scales down the gradients proportionally if their norm is
228/// greater than `max_norm`.
229pub fn clip_grad_norm(
230    gradients: &FlattenedModuleParam,
231    max_norm: f32,
232) -> crate::error::Result<(MaybeClippedGrads<'_>, f32)> {
233    let total_norm: f32 = gradients
234        .values()
235        .try_fold(array!(0.0), |acc, grad| acc.add(&grad.square()?.sum(None)?))?
236        .sqrt()?
237        .item_exact();
238    let normalizer = array!(max_norm / (total_norm + 1e-6));
239
240    let clipped_gradients: HashMap<_, _> = gradients
241        .iter()
242        .map(|(key, grad)| {
243            let clipped_grad = if total_norm < max_norm {
244                Cow::Borrowed(grad)
245            } else {
246                Cow::Owned(grad * &normalizer)
247            };
248            (key.clone(), clipped_grad)
249        })
250        .collect();
251    Ok((clipped_gradients, total_norm))
252}
253
254#[cfg(test)]
255mod tests {
256    use std::collections::HashMap;
257
258    use crate::{
259        array,
260        module::FlattenedModuleParam,
261        test_utils::{assert_array_eq, tolerances},
262        Array,
263    };
264
265    use super::clip_grad_norm;
266
267    #[test]
268    fn test_clip_grad_norm() {
269        // Test with small gradients that do not require clipping
270        let mut small_grads: FlattenedModuleParam = HashMap::new();
271        small_grads.insert("first.a".into(), array!([0.1, 0.2]));
272        small_grads.insert("first.b".into(), array!(0.1));
273        small_grads.insert("second".into(), array!(0.3));
274
275        let max_norm = 10.0;
276
277        let (clipped_grads, _) = clip_grad_norm(&small_grads, max_norm).unwrap();
278        for (key, value) in small_grads.iter() {
279            assert_array_eq(
280                &*clipped_grads[key],
281                value,
282                tolerances::EXACT.rtol,
283                tolerances::EXACT.atol,
284            );
285        }
286
287        // Test with large gradients that require clipping
288        let mut large_grads: FlattenedModuleParam = HashMap::new();
289        large_grads.insert("first.a".into(), array!([10.0, 20.0]));
290        large_grads.insert("first.b".into(), array!(10.0));
291        large_grads.insert("second".into(), array!(30.0));
292
293        let max_norm = 1.0;
294
295        let (clipped_grads, total_norm) = clip_grad_norm(&large_grads, max_norm).unwrap();
296        let clipped_values: Vec<_> = clipped_grads.values().map(|v| v.as_ref()).collect();
297        let norm_of_clipped = clipped_values
298            .into_iter()
299            .map(|g| g.square().unwrap().sum(None).unwrap())
300            .sum::<Array>()
301            .sqrt()
302            .unwrap();
303
304        float_eq::assert_float_eq!(norm_of_clipped.item_exact::<f32>(), max_norm, abs <= 1e-6);
305
306        // Ensures that the scaling was done correctly
307        let scale = max_norm / total_norm;
308        let expected_grads: FlattenedModuleParam = large_grads
309            .iter()
310            .map(|(key, value)| (key.clone(), value * scale))
311            .collect();
312        for (key, value) in expected_grads.iter() {
313            assert_array_eq(
314                &*clipped_grads[key],
315                value,
316                tolerances::EXACT.rtol,
317                tolerances::EXACT.atol,
318            );
319        }
320    }
321}
mlx_rs/optimizers/mod.rs

mlx_rs/optimizers/
mod.rs
