rlx-optim 0.2.18

// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
// SPDX-License-Identifier: MIT OR Apache-2.0

//! Stochastic Gradient Descent with optional momentum and decoupled
//! L2 weight decay.
//!
//! # Update rules
//!
//! Vanilla SGD (`momentum = 0`):
//!
//! ```text
//! θ_{t+1} = θ_t − lr · (g_t + λ·θ_t)
//! ```
//!
//! Polyak momentum (`momentum = μ`, `nesterov = false`):
//!
//! ```text
//! v_{t+1} = μ·v_t + (g_t + λ·θ_t)
//! θ_{t+1} = θ_t − lr · v_{t+1}
//! ```
//!
//! Nesterov-accelerated SGD (`nesterov = true`):
//!
//! ```text
//! v_{t+1} = μ·v_t + (g_t + λ·θ_t)
//! θ_{t+1} = θ_t − lr · (g_t + λ·θ_t + μ·v_{t+1})
//! ```
//!
//! # When to use
//!
//! The default choice when training CNNs from scratch; with a
//! well-tuned `lr` schedule it still beats Adam on many vision
//! benchmarks. Cheap state (one buffer if `momentum > 0`).

use std::collections::HashMap;

use crate::Optimizer;
use crate::common::zeros_entry;

/// SGD with momentum / Nesterov / L2 weight decay.
///
/// All hyperparameters are public so callers can hot-swap them between
/// iterations (e.g. for a warm-up schedule). State is keyed by
/// parameter name; the same `Sgd` instance can drive every tensor in
/// a model.
#[derive(Debug, Clone)]
pub struct Sgd {
    /// Learning rate. No default — pass it to [`Sgd::new`].
    pub lr: f32,
    /// Polyak momentum coefficient ∈ \[0, 1\). `0.0` disables momentum
    /// entirely (and the per-tensor velocity buffer is still allocated
    /// but unused — set via [`Sgd::with_momentum`] if you want it on).
    pub momentum: f32,
    /// Use Nesterov-accelerated momentum. Only meaningful when
    /// `momentum > 0`.
    pub nesterov: bool,
    /// L2 weight decay coefficient λ. Folded into the gradient
    /// *before* the momentum EMA (classical, **not** decoupled).
    /// Use [`crate::AdamW`]-style decoupling if you need that.
    pub weight_decay: f32,
    v: HashMap<String, Vec<f32>>,
}

impl Sgd {
    /// Construct with `lr` and momentum / decay disabled.
    pub fn new(lr: f32) -> Self {
        Self {
            lr,
            momentum: 0.0,
            nesterov: false,
            weight_decay: 0.0,
            v: HashMap::new(),
        }
    }

    /// Enable Polyak (or Nesterov) momentum.
    pub fn with_momentum(mut self, momentum: f32, nesterov: bool) -> Self {
        self.momentum = momentum;
        self.nesterov = nesterov;
        self
    }

    /// Set the L2 weight-decay coefficient.
    pub fn with_weight_decay(mut self, wd: f32) -> Self {
        self.weight_decay = wd;
        self
    }
}

// ── checkpointing ───────────────────────────────────────────

impl Sgd {
    /// Named accumulators plus the step counter — see
    /// [`crate::OptimizerState`].
    pub(crate) fn snapshot(&self) -> crate::OptimizerState {
        let mut out = crate::OptimizerState {
            step: 0,
            buffers: Vec::new(),
        };
        out.extend_slot("v", &self.v);
        out
    }

    pub(crate) fn restore(&mut self, state: &crate::OptimizerState) {
        state.take_slot("v", &mut self.v);
    }
}

impl Optimizer for Sgd {
    fn set_lr(&mut self, lr: f32) {
        self.lr = lr;
    }

    fn state_dict(&self) -> Option<crate::OptimizerState> {
        Some(self.snapshot())
    }

    fn load_state_dict(&mut self, state: &crate::OptimizerState) -> bool {
        self.restore(state);
        true
    }

    fn step(&mut self, name: &str, _shape: &[usize], param: &mut [f32], grad: &[f32]) {
        debug_assert_eq!(param.len(), grad.len());
        let mu = self.momentum;
        let wd = self.weight_decay;
        let lr = self.lr;
        // Read out of `self` *before* the loop. Both of these used to be tested
        // per element, and `self.nesterov` was read through a live `&mut self`
        // borrow, which stopped the loop vectorizing: SGD measured 1.06
        // ns/element against Lion's 0.28 while doing strictly less arithmetic.
        let nesterov = self.nesterov;
        let v = zeros_entry(&mut self.v, name, param.len());

        // One specialized loop per configuration rather than one loop with the
        // configuration inside it. The arithmetic is unchanged, so results stay
        // bit-identical.
        if mu == 0.0 {
            for (p, g) in param.iter_mut().zip(grad) {
                let d = *g + wd * *p;
                *p -= lr * d;
            }
        } else if nesterov {
            for ((p, vi), g) in param.iter_mut().zip(v.iter_mut()).zip(grad) {
                let d = *g + wd * *p;
                *vi = mu * *vi + d;
                *p -= lr * (d + mu * *vi);
            }
        } else {
            for ((p, vi), g) in param.iter_mut().zip(v.iter_mut()).zip(grad) {
                let d = *g + wd * *p;
                *vi = mu * *vi + d;
                *p -= lr * *vi;
            }
        }
    }
}