forge/optim/
mod.rs

1//! Optimizers (roadmap v4, Stage 9): AdamW with decoupled weight decay and
2//! optional global gradient-norm clipping. Backend-agnostic — state lives on
3//! the same device as the parameters and updates run through `ops`.
4
5use crate::error::{ForgeError, Result};
6use crate::ops;
7use crate::tensor::Tensor;
8
9#[derive(Clone, Copy, Debug)]
10pub struct AdamWOpts {
11    pub lr: f32,
12    pub beta1: f32,
13    pub beta2: f32,
14    pub eps: f32,
15    pub weight_decay: f32,
16    /// Clip the global gradient norm to this value when set.
17    pub clip: Option<f32>,
18}
19
20impl Default for AdamWOpts {
21    fn default() -> Self {
22        AdamWOpts {
23            lr: 3e-4,
24            beta1: 0.9,
25            beta2: 0.95,
26            eps: 1e-8,
27            weight_decay: 0.1,
28            clip: Some(1.0),
29        }
30    }
31}
32
33pub struct AdamW {
34    pub opts: AdamWOpts,
35    step: u32,
36    /// (m, v) per parameter, in registration order.
37    state: Vec<(Tensor, Tensor)>,
38    /// Whether weight decay applies (false for biases and LayerNorm params).
39    decay: Vec<bool>,
40}
41
42impl AdamW {
43    /// `params` supplies each parameter with its decay flag; order must
44    /// match the `step` calls.
45    pub fn new(params: &[(&Tensor, bool)], opts: AdamWOpts) -> Result<AdamW> {
46        let mut state = Vec::with_capacity(params.len());
47        let mut decay = Vec::with_capacity(params.len());
48        for (p, d) in params {
49            let device = p.device();
50            state.push((
51                Tensor::zeros(p.shape().clone(), &device)?,
52                Tensor::zeros(p.shape().clone(), &device)?,
53            ));
54            decay.push(*d);
55        }
56        Ok(AdamW {
57            opts,
58            step: 0,
59            state,
60            decay,
61        })
62    }
63
64    /// Apply one update. Returns the pre-clip global gradient norm.
65    pub fn step(&mut self, params: &mut [&mut Tensor], grads: &[Tensor]) -> Result<f32> {
66        if params.len() != self.state.len() || grads.len() != self.state.len() {
67            return Err(ForgeError::Shape(format!(
68                "AdamW::step arity mismatch: {} params, {} grads, {} slots",
69                params.len(),
70                grads.len(),
71                self.state.len()
72            )));
73        }
74        self.step += 1;
75        let mut total_sq = 0.0f32;
76        for g in grads {
77            total_sq += ops::sumsq(g)?;
78        }
79        let norm = total_sq.sqrt();
80        let rescale = match self.opts.clip {
81            Some(c) if norm > c && norm > 0.0 => Some(c / norm),
82            _ => None,
83        };
84        let mut clipped;
85        for ((p, g), ((m, v), decay)) in params
86            .iter_mut()
87            .zip(grads)
88            .zip(self.state.iter_mut().zip(&self.decay))
89        {
90            let g = match rescale {
91                Some(s) => {
92                    clipped = ops::scale(g, s)?;
93                    &clipped
94                }
95                None => g,
96            };
97            ops::adamw(
98                p,
99                g,
100                m,
101                v,
102                self.opts.lr,
103                self.opts.beta1,
104                self.opts.beta2,
105                self.opts.eps,
106                if *decay { self.opts.weight_decay } else { 0.0 },
107                self.step,
108            )?;
109        }
110        Ok(norm)
111    }
112}
forge/optim/mod.rs

forge/optim/
mod.rs
