rlx_optim/
lib.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5//! RLX training-step optimizers.
6//!
7//! Host-side `f32` step functions for the families surveyed in
8//! "A Systematic Review of Optimization Algorithms for Modern Deep
9//! Learning" (arXiv:2509.02046v1). Each algorithm exposes a small
10//! state struct keyed by parameter *name* (so the same struct holds
11//! moments for every tensor in a model) and a `step` method that
12//! consumes `(name, shape, &mut params, &grads)`.
13//!
14//! The API is deliberately minimal: it operates on flat `&mut [f32]`
15//! / `&[f32]` slices plus a `&[usize]` shape — matching the
16//! [`rlx_umap::adam`](../rlx_umap/adam/index.html) pattern. Backends
17//! that already ship a fused step kernel (see e.g.
18//! `rlx_metal::splat_adam`) are free to bypass this crate for their
19//! hot path; this crate is the portable reference / CPU fallback / the
20//! one used when there is no backend fused kernel for the requested
21//! algorithm.
22//!
23//! # Algorithms
24//!
25//! | Family          | Type                          |
26//! |-----------------|-------------------------------|
27//! | [`Sgd`]         | SGD ± momentum / Nesterov     |
28//! | [`Adam`]        | Adam                          |
29//! | [`AdamW`]       | AdamW (decoupled decay)       |
30//! | [`NAdamW`]      | Nesterov AdamW                |
31//! | [`RAdam`]       | Rectified Adam                |
32//! | [`QHAdamW`]     | Quasi-hyperbolic AdamW        |
33//! | [`Lamb`]        | LAMB (layer-wise adaptive)    |
34//! | [`Adafactor`]   | Adafactor (factored 2nd mom.) |
35//! | [`Lion`]        | Lion (sign of EMA)            |
36//! | [`Soap`]        | SOAP (Shampoo-in-Adam-basis)  |
37//! | [`KronPsgd`]    | Kron / PSGD                   |
38//! | [`Muon`]        | Muon (Newton–Schulz orth.)    |
39//! | [`Sophia`]      | Sophia-H                      |
40//! | [`Mars`]        | MARS (variance-reduced)       |
41//! | [`Stiefel`]     | Riemannian SGD on St(m,n)     |
42
43// Pure-safe by default. Relaxed from `forbid` to `deny` so the one
44// performance-critical exception — the Accelerate/AMX `cblas_sgemm` shim behind
45// Muon's Newton–Schulz on macOS — can opt in with a scoped `#[allow(unsafe_code)]`
46// (see `muon::accel`). Every other module stays unsafe-free.
47#![deny(unsafe_code)]
48
49mod common;
50
51mod adafactor;
52mod adam;
53mod adamw;
54mod kron_psgd;
55mod lamb;
56mod lion;
57mod mars;
58mod muon;
59mod nadamw;
60mod qhadamw;
61mod radam;
62mod sgd;
63mod soap;
64mod sophia;
65mod stiefel;
66
67pub use adafactor::Adafactor;
68pub use adam::Adam;
69pub use adamw::AdamW;
70pub use kron_psgd::KronPsgd;
71pub use lamb::Lamb;
72pub use lion::Lion;
73pub use mars::Mars;
74pub use muon::{Muon, newton_schulz_orth};
75pub use nadamw::NAdamW;
76pub use qhadamw::QHAdamW;
77pub use radam::RAdam;
78pub use sgd::Sgd;
79pub use soap::Soap;
80pub use sophia::Sophia;
81pub use stiefel::Stiefel;
82
83pub use common::{global_grad_clip_scale, l2_norm};
84
85/// Common parameter-update interface.
86///
87/// `name` keys the per-parameter state (moments, preconditioners),
88/// `shape` is the parameter's logical shape (used by matrix-aware
89/// algorithms like Adafactor / SOAP / Muon — ignored by elementwise
90/// ones), `param` is updated in place from `grad`. `grad` is treated
91/// as read-only; callers that need gradient clipping should pre-scale
92/// it (see [`global_grad_clip_scale`]).
93///
94/// # Implementing for a backend
95///
96/// Every algorithm in this crate provides a CPU reference impl. A
97/// backend (e.g. `rlx-metal`, `rlx-cuda`) is free to write its own
98/// fused step kernel and impl `Optimizer` for a wrapper struct that
99/// owns device buffers — the trait places no requirement on where
100/// the state lives, only on the entry-point signature. The
101/// `rlx-metal::splat_adam` kernel is the canonical example of a
102/// backend that bypasses this crate entirely; you can wrap it with a
103/// 5-line `impl Optimizer` if you want a uniform interface from a
104/// generic trainer.
105///
106/// # Per-tensor learning rate
107///
108/// For optimizers that don't need per-tensor LR variation (most
109/// transformer pre-training), set [`lr_scale`](Self::lr_scale) to
110/// return `1.0` (the default). For domain-specific use cases — e.g.
111/// 3D Gaussian splatting, where different attributes need wildly
112/// different step sizes — override [`lr_scale`](Self::lr_scale) to
113/// multiply the base `lr` by a per-name factor. The provided method
114/// on the trait does NOT scale automatically; algorithms are free to
115/// consult it via [`Optimizer::lr_scale`] inside their `step`.
116/// One parameter's data for a batched optimizer step ([`Optimizer::step_batch`]):
117/// its name, static shape, mutable data slice (owned by the caller), and gradient.
118pub struct OptItem<'a> {
119    pub name: &'a str,
120    pub shape: &'a [usize],
121    pub param: &'a mut [f32],
122    pub grad: &'a [f32],
123}
124
125pub trait Optimizer {
126    fn step(&mut self, name: &str, shape: &[usize], param: &mut [f32], grad: &[f32]);
127
128    /// Batched step over ALL parameters in one call. Default: sequential
129    /// `step` per item — bit-identical to the per-parameter loop. Optimizers
130    /// whose parameter groups are **independent** (e.g. Muon on the 2-D weight
131    /// matrices vs AdamW on the embeddings/biases/norms) can override this to
132    /// run the groups on separate threads; because the groups touch disjoint
133    /// parameters and disjoint optimizer state, the result is bit-for-bit the
134    /// same as the serial loop — only the wall-clock (the CPU-side optimizer
135    /// bubble) shrinks toward `max(group_times)` instead of their sum.
136    fn step_batch(&mut self, items: &mut [OptItem<'_>]) {
137        for it in items.iter_mut() {
138            self.step(it.name, it.shape, it.param, it.grad);
139        }
140    }
141
142    /// Advance the global step counter. Most algorithms increment per
143    /// call to [`step`], so most implementations leave this a no-op.
144    fn end_iteration(&mut self) {}
145
146    /// Set the base learning rate (for LR schedules / warmup). Default is a
147    /// no-op for algorithms without a scalar `lr` (e.g. Adafactor's relative
148    /// step size); every algorithm in this crate that has an `lr` field
149    /// overrides this to update it.
150    fn set_lr(&mut self, _lr: f32) {}
151
152    /// Per-tensor multiplier on the effective learning rate. Default
153    /// is `1.0` for every name. Override when wrapping this crate to
154    /// support per-name LR schedules (e.g. embedding-vs-attention
155    /// splits, or the Gaussian-splat attribute-typed LR setup). The
156    /// CPU impls in this crate currently honor this only when the
157    /// caller passes a pre-scaled `lr` for the relevant call —
158    /// backends are encouraged to consult it inside their fused
159    /// kernel.
160    fn lr_scale(&self, _name: &str) -> f32 {
161        1.0
162    }
163}
rlx_optim/lib.rs

rlx_optim/
lib.rs
