burn_optim/optim/module/
base.rs

1use alloc::sync::Arc;
2use core::any::Any;
3
4use crate::{RecordState, StateSink, StateSource};
5use burn_core as burn;
6use burn_core::tensor::kind::BridgeTensor;
7
8use crate::LearningRate;
9use burn::tensor::{Device, Tensor};
10
11/// An opinionated trait to simplify the process of implementing an optimizer.
12///
13/// Implementations don't have to handle missing gradients, loading and exporting records,
14/// navigate the module parameter structure, handle tracked and untracked tensors, and the likes.
15/// Wrap one in a [`ModuleOptimizer`](crate::optim::ModuleOptimizer) to optimize a whole module.
16pub trait Optimizer: Send + Sync + Clone + 'static {
17    /// The state of the optimizer for a single parameter of rank `D`.
18    ///
19    /// It implements [`RecordState`] (which itself requires `Send + Sync + 'static`) so it can be
20    /// decomposed into named tensors and scalars for the burnpack format.
21    type State<const D: usize>: Clone + RecordState;
22
23    /// The optimizer step is performed for one tensor at a time with its gradient and state.
24    ///
25    /// Note that the state is passed as parameter, so implementations don't have to handle
26    /// the saving and loading of recorded states.
27    fn step<const D: usize>(
28        &self,
29        lr: LearningRate,
30        tensor: Tensor<D>,
31        grad: Tensor<D>,
32        state: Option<Self::State<D>>,
33    ) -> (Tensor<D>, Option<Self::State<D>>);
34
35    /// Change the device of the state.
36    ///
37    /// This function will be called accordingly to have the state on the same device as the
38    /// gradient and the tensor when the [step](Optimizer::step) function is called.
39    fn to_device<const D: usize>(state: Self::State<D>, device: &Device) -> Self::State<D>;
40}
41
42/// A type-erased optimizer state for a single parameter.
43///
44/// It wraps a concrete `O::State<D>` together with its rank `D` so that the rank can be recovered
45/// when the state is later interpreted by the originating [`Optimizer`] (during a step, a device
46/// transfer or serialization).
47#[derive(Clone)]
48pub struct DynState {
49    state: Arc<dyn Any + Send + Sync>,
50    rank: usize,
51}
52
53impl DynState {
54    /// Erase a concrete optimizer state of rank `rank`.
55    pub fn create<T: Send + Sync + 'static>(state: T, rank: usize) -> Self {
56        Self {
57            state: Arc::new(state),
58            rank,
59        }
60    }
61
62    /// Recover the concrete state by value, moving it out when this is the only handle and
63    /// cloning only when the underlying state is still shared. Panics if `T` does not match the
64    /// stored type.
65    pub fn downcast<T: Clone + Send + Sync + 'static>(self) -> T {
66        let state = self
67            .state
68            .downcast::<T>()
69            .expect("The dynamic optimizer state should match the optimizer state type.");
70        Arc::try_unwrap(state).unwrap_or_else(|state| (*state).clone())
71    }
72
73    /// Borrow the concrete state without cloning. Panics if `T` does not match the stored type.
74    pub fn downcast_ref<T: 'static>(&self) -> &T {
75        self.state
76            .downcast_ref::<T>()
77            .expect("The dynamic optimizer state should match the optimizer state type.")
78    }
79
80    /// The rank of the parameter this state belongs to.
81    pub fn rank(&self) -> usize {
82        self.rank
83    }
84}
85
86/// Dispatch a runtime `rank` to a body parameterized by a `const D: usize`.
87macro_rules! dispatch_rank {
88    ($rank:expr, $d:ident => $body:block) => {
89        match $rank {
90            0 => {
91                const $d: usize = 0;
92                $body
93            }
94            1 => {
95                const $d: usize = 1;
96                $body
97            }
98            2 => {
99                const $d: usize = 2;
100                $body
101            }
102            3 => {
103                const $d: usize = 3;
104                $body
105            }
106            4 => {
107                const $d: usize = 4;
108                $body
109            }
110            5 => {
111                const $d: usize = 5;
112                $body
113            }
114            6 => {
115                const $d: usize = 6;
116                $body
117            }
118            7 => {
119                const $d: usize = 7;
120                $body
121            }
122            8 => {
123                const $d: usize = 8;
124                $body
125            }
126            other => panic!("Unsupported tensor rank for optimizer state: {other}"),
127        }
128    };
129}
130
131/// Object-safe view over an [`Optimizer`], allowing [`ModuleOptimizer`](crate::optim::ModuleOptimizer)
132/// to stay non-generic. Rank-generic operations are dispatched on a runtime rank.
133pub trait DynOptimizer: Send + Sync {
134    /// Perform an optimizer step for a single parameter of the given `rank`.
135    fn step_dyn(
136        &self,
137        rank: usize,
138        lr: LearningRate,
139        tensor: BridgeTensor,
140        grad: BridgeTensor,
141        state: Option<DynState>,
142    ) -> (BridgeTensor, Option<DynState>);
143
144    /// Move a state to the given device.
145    fn to_device_dyn(&self, state: DynState, device: &Device) -> DynState;
146
147    /// Decompose a state into named tensors and scalars under `prefix`.
148    fn state_flatten(&self, prefix: &str, state: &DynState, out: &mut StateSink);
149
150    /// Rebuild a state of the given `rank` from named tensors and scalars under `prefix`.
151    ///
152    /// Returns `None` when the record does not contain a reconstructable state for this parameter
153    /// (e.g. a truncated or foreign file); the caller leaves that parameter without state, so it is
154    /// re-initialized on the next step.
155    fn state_unflatten(
156        &self,
157        rank: usize,
158        prefix: &str,
159        src: &mut StateSource,
160        device: &Device,
161    ) -> Option<DynState>;
162}
163
164impl<O: Optimizer> DynOptimizer for O {
165    fn step_dyn(
166        &self,
167        rank: usize,
168        lr: LearningRate,
169        tensor: BridgeTensor,
170        grad: BridgeTensor,
171        state: Option<DynState>,
172    ) -> (BridgeTensor, Option<DynState>) {
173        dispatch_rank!(rank, D => {
174            let (tensor, state) = self.step(
175                lr,
176                Tensor::<D>::from_bridge(tensor),
177                Tensor::<D>::from_bridge(grad),
178                state.map(|state| state.downcast::<O::State<D>>()),
179            );
180
181            (tensor.into_bridge(), state.map(|state| DynState::create(state, D)))
182        })
183    }
184
185    fn to_device_dyn(&self, state: DynState, device: &Device) -> DynState {
186        dispatch_rank!(state.rank(), D => {
187            let state = O::to_device::<D>(state.downcast::<O::State<D>>(), device);
188            DynState::create(state, D)
189        })
190    }
191
192    fn state_flatten(&self, prefix: &str, state: &DynState, out: &mut StateSink) {
193        dispatch_rank!(state.rank(), D => {
194            RecordState::state_flatten(state.downcast_ref::<O::State<D>>(), prefix, out);
195        })
196    }
197
198    fn state_unflatten(
199        &self,
200        rank: usize,
201        prefix: &str,
202        src: &mut StateSource,
203        device: &Device,
204    ) -> Option<DynState> {
205        dispatch_rank!(rank, D => {
206            let state = <O::State<D> as RecordState>::state_unflatten(prefix, src, device)?;
207            Some(DynState::create(state, D))
208        })
209    }
210}
burn_optim/optim/module/base.rs

burn_optim/optim/module/
base.rs
