feat(batch14): Muon optimizer, logit processors, per-token activation quantization

- MuonOptimizer: Nesterov momentum + quintic Newton-Schulz orthogonalization (arXiv:2409.20325); 5-iteration NS maps gradient to near-orthogonal matrix; 1D fallback skips NS; decoupled weight decay; 16 tests - LogitProcessorList: temperature, top-k, top-p nucleus, min-p, repetition/ presence/frequency penalty, eta-sampling; softmax/log_softmax/argmax/ sample_token helpers; 39 tests - ActivationQuantizer: per-token dynamic INT8/FP8E4M3 scaling for inference activations; per-tensor mode; dequantize; max_error diagnostic; 19 tests Co-Authored-By: Claude Sonnet 4.6 <[email protected]>
2026-06-27 06:11:25 +00:00
parent e2d902b34b
commit e00a6da018
6 changed files with 2016 additions and 0 deletions
@@ -0,0 +1,534 @@
+//! Muon optimizer — Nesterov momentum with Newton-Schulz orthogonalization.
+//!
+//! Implements the Muon algorithm (Jordan et al., arXiv:2409.20325, 2024):
+//! weight-matrix updates are orthogonalized via a quintic Newton-Schulz
+//! iteration before being applied, yielding near-unit spectral norm updates
+//! and significantly outperforming AdamW on language modelling benchmarks.
+//!
+//! # Algorithm
+//!
+//! For each parameter matrix W of shape \[m × n\]:
+//!
+//! 1. **Nesterov momentum**: `m_t = β·m_{t-1} + g_t`;  `ĝ = g_t + β·m_t`
+//! 2. **Newton-Schulz orthogonalization** (5 iterations, quintic polynomial,
+//!    only for 2-D matrices): maps `ĝ` to a near-orthogonal matrix with
+//!    spectral norm ≈ 1.
+//! 3. **Parameter update**: `W ← (1 − λ·lr)·W − lr·O`
+//!
+//! 1-D parameters (bias vectors, embedding rows) fall back to plain SGD
+//! without orthogonalization.
+//!
+//! # Example
+//!
+//! ```rust
+//! use rtx_transformers::optimizers::muon::{MuonConfig, MuonOptimizer};
+//!
+//! let mut opt = MuonOptimizer::new(MuonConfig::default());
+//! opt.register("W", 4, 4);
+//!
+//! let params: Vec<f32> = vec![1.0; 16];
+//! let grad: Vec<f32>   = vec![0.1; 16];
+//! let updated = opt.step("W", &params, &grad);
+//! assert_eq!(updated.len(), 16);
+//! ```
+
+use std::collections::HashMap;
+
+// ============================================================================
+// Configuration
+// ============================================================================
+
+/// Configuration for the Muon optimizer.
+#[derive(Debug, Clone)]
+pub struct MuonConfig {
+    /// Learning rate (default 0.02).
+    pub lr: f32,
+    /// Nesterov momentum coefficient β (default 0.95).
+    pub momentum: f32,
+    /// Small constant added to Frobenius norm for numerical stability (default 1e-8).
+    pub epsilon: f32,
+    /// Number of Newton-Schulz iterations (default 5).
+    pub ns_steps: usize,
+    /// Decoupled weight-decay coefficient λ (default 0.0 — disabled).
+    pub weight_decay: f32,
+}
+
+impl Default for MuonConfig {
+    fn default() -> Self {
+        Self {
+            lr: 0.02,
+            momentum: 0.95,
+            epsilon: 1e-8,
+            ns_steps: 5,
+            weight_decay: 0.0,
+        }
+    }
+}
+
+// ============================================================================
+// Per-parameter state
+// ============================================================================
+
+/// Optimizer state stored for a single parameter.
+pub struct MuonParamState {
+    /// Momentum buffer (flattened, same length as `rows * cols`).
+    pub momentum_buffer: Vec<f32>,
+    /// Number of rows of the parameter matrix.
+    pub rows: usize,
+    /// Number of columns of the parameter matrix.
+    /// When `cols == 1` the parameter is treated as 1-D (no orthogonalization).
+    pub cols: usize,
+}
+
+// ============================================================================
+// Optimizer
+// ============================================================================
+
+/// Muon optimizer — CPU reference implementation.
+pub struct MuonOptimizer {
+    config: MuonConfig,
+    states: HashMap<String, MuonParamState>,
+}
+
+impl MuonOptimizer {
+    /// Create a new `MuonOptimizer` with the given configuration.
+    #[must_use]
+    pub fn new(config: MuonConfig) -> Self {
+        Self {
+            config,
+            states: HashMap::new(),
+        }
+    }
+
+    /// Convenience constructor: default config with `lr` overridden.
+    #[must_use]
+    pub fn with_lr(lr: f32) -> Self {
+        Self::new(MuonConfig { lr, ..MuonConfig::default() })
+    }
+
+    /// Register a parameter before the first `step()` call.
+    ///
+    /// `rows` and `cols` describe the logical matrix shape of the (flattened)
+    /// parameter.  Set `cols = 1` to request the 1-D (plain-SGD) code path.
+    pub fn register(&mut self, name: &str, rows: usize, cols: usize) {
+        let n = rows * cols;
+        self.states.insert(
+            name.to_owned(),
+            MuonParamState {
+                momentum_buffer: vec![0.0_f32; n],
+                rows,
+                cols,
+            },
+        );
+    }
+
+    // -----------------------------------------------------------------------
+    // Public step
+    // -----------------------------------------------------------------------
+
+    /// Perform one optimizer step for the named parameter.
+    ///
+    /// # Panics
+    ///
+    /// Panics if `name` was not previously registered with [`Self::register`].
+    pub fn step(&mut self, name: &str, params: &[f32], grad: &[f32]) -> Vec<f32> {
+        // Borrow config values we need so we do not borrow `self` mutably and
+        // immutably at the same time.
+        let momentum = self.config.momentum;
+        let lr = self.config.lr;
+        let weight_decay = self.config.weight_decay;
+
+        let state = self
+            .states
+            .get_mut(name)
+            .unwrap_or_else(|| panic!("MuonOptimizer: parameter '{name}' was not registered"));
+
+        let rows = state.rows;
+        let cols = state.cols;
+
+        // 1. Nesterov momentum accumulation
+        //    m_t = β·m_{t-1} + g_t
+        for (m, g) in state.momentum_buffer.iter_mut().zip(grad.iter()) {
+            *m = momentum * *m + g;
+        }
+
+        //    ĝ = g_t + β·m_t
+        let nesterov: Vec<f32> = grad
+            .iter()
+            .zip(state.momentum_buffer.iter())
+            .map(|(g, m)| g + momentum * m)
+            .collect();
+
+        // 2. Orthogonalize if 2-D (cols > 1), else pass through
+        let update = if cols > 1 {
+            self.newton_schulz5(&nesterov, rows, cols)
+        } else {
+            nesterov
+        };
+
+        // 3. Parameter update with optional decoupled weight decay
+        //    p ← (1 − λ·lr)·p − lr·u
+        params
+            .iter()
+            .zip(update.iter())
+            .map(|(p, u)| (1.0 - weight_decay * lr) * p - lr * u)
+            .collect()
+    }
+
+    // -----------------------------------------------------------------------
+    // Newton-Schulz orthogonalization
+    // -----------------------------------------------------------------------
+
+    /// Quintic Newton-Schulz iteration mapping `g` (shape `rows × cols`) to a
+    /// near-orthogonal matrix whose Frobenius norm is ≤ `max(rows, cols)^{1/2}`.
+    ///
+    /// The spectral norm of the output converges to ≈ 1 after `ns_steps`
+    /// iterations, making the update scale-independent and well-conditioned.
+    ///
+    /// Coefficients `(a, b, c) = (3.4445, −4.7750, 2.0315)` are those of
+    /// Kosson et al. (2023) quintic polynomial.
+    #[must_use]
+    pub fn newton_schulz5(&self, g: &[f32], rows: usize, cols: usize) -> Vec<f32> {
+        // Normalize to Frobenius norm ≈ 1 for numerically stable iteration
+        let norm = Self::frobenius_norm(g) + self.config.epsilon;
+        let mut x: Vec<f32> = g.iter().map(|v| v / norm).collect();
+
+        // Quintic coefficients (Kosson et al.)
+        let a = 3.4445_f32;
+        let b = -4.7750_f32;
+        let c = 2.0315_f32;
+
+        for _ in 0..self.config.ns_steps {
+            // X^T  [cols × rows]
+            let xt = Self::transpose(&x, rows, cols);
+
+            // X^T · X  [cols × cols]
+            let xt_x = Self::matmul(&xt, &x, cols, rows, cols);
+
+            // (X^T · X)²  [cols × cols]
+            let xt_x2 = Self::matmul(&xt_x, &xt_x, cols, cols, cols);
+
+            // X · (X^T · X)  [rows × cols]
+            let x_xt_x = Self::matmul(&x, &xt_x, rows, cols, cols);
+
+            // X · (X^T · X)²  [rows × cols]
+            let x_xt_x2 = Self::matmul(&x, &xt_x2, rows, cols, cols);
+
+            // X_{k+1} = a·X + b·(X·X^T·X) + c·(X·(X^T·X)²)
+            x = x
+                .iter()
+                .zip(x_xt_x.iter())
+                .zip(x_xt_x2.iter())
+                .map(|((xi, bxi), cxi)| a * xi + b * bxi + c * cxi)
+                .collect();
+        }
+
+        x
+    }
+
+    // -----------------------------------------------------------------------
+    // Linear-algebra helpers (row-major, plain CPU)
+    // -----------------------------------------------------------------------
+
+    /// Dense matrix multiply: `A[m × k] @ B[k × n] → C[m × n]`.
+    fn matmul(a: &[f32], b: &[f32], m: usize, k: usize, n: usize) -> Vec<f32> {
+        debug_assert_eq!(a.len(), m * k);
+        debug_assert_eq!(b.len(), k * n);
+
+        let mut c = vec![0.0_f32; m * n];
+        for i in 0..m {
+            for p in 0..k {
+                let a_ip = a[i * k + p];
+                for j in 0..n {
+                    c[i * n + j] += a_ip * b[p * n + j];
+                }
+            }
+        }
+        c
+    }
+
+    /// Transpose `A[m × n]` → `A^T[n × m]`.
+    fn transpose(a: &[f32], m: usize, n: usize) -> Vec<f32> {
+        debug_assert_eq!(a.len(), m * n);
+
+        let mut at = vec![0.0_f32; n * m];
+        for i in 0..m {
+            for j in 0..n {
+                at[j * m + i] = a[i * n + j];
+            }
+        }
+        at
+    }
+
+    /// Frobenius norm: `||A||_F = sqrt(Σ aᵢⱼ²)`.
+    fn frobenius_norm(a: &[f32]) -> f32 {
+        a.iter().map(|x| x * x).sum::<f32>().sqrt()
+    }
+}
+
+// ============================================================================
+// Tests
+// ============================================================================
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    // -------------------------------------------------------------------------
+    // Construction
+    // -------------------------------------------------------------------------
+
+    #[test]
+    fn test_muon_creation() {
+        let cfg = MuonConfig::default();
+        assert!((cfg.lr - 0.02).abs() < 1e-9, "default lr should be 0.02");
+        assert!((cfg.momentum - 0.95).abs() < 1e-9, "default momentum should be 0.95");
+        assert!((cfg.epsilon - 1e-8).abs() < 1e-15, "default epsilon should be 1e-8");
+        assert_eq!(cfg.ns_steps, 5, "default ns_steps should be 5");
+        assert!((cfg.weight_decay - 0.0).abs() < 1e-9, "default weight_decay should be 0.0");
+    }
+
+    #[test]
+    fn test_with_lr() {
+        let opt = MuonOptimizer::with_lr(0.001);
+        assert!((opt.config.lr - 0.001).abs() < 1e-9);
+        // All other fields stay at their defaults
+        assert!((opt.config.momentum - 0.95).abs() < 1e-9);
+        assert_eq!(opt.config.ns_steps, 5);
+    }
+
+    // -------------------------------------------------------------------------
+    // Math primitives
+    // -------------------------------------------------------------------------
+
+    #[test]
+    fn test_frobenius_norm() {
+        // ||(3, 4)||_F = 5.0
+        let v = vec![3.0_f32, 4.0];
+        let n = MuonOptimizer::frobenius_norm(&v);
+        assert!((n - 5.0).abs() < 1e-6, "expected 5.0 got {n}");
+    }
+
+    #[test]
+    fn test_transpose_2x3() {
+        // A = [[1,2,3],[4,5,6]]  →  A^T = [[1,4],[2,5],[3,6]]
+        let a = vec![1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0];
+        let at = MuonOptimizer::transpose(&a, 2, 3);
+        let expected = vec![1.0_f32, 4.0, 2.0, 5.0, 3.0, 6.0];
+        assert_eq!(at, expected, "transpose mismatch");
+    }
+
+    #[test]
+    fn test_matmul_identity() {
+        // A @ I == A  for a 3×3 matrix
+        let a: Vec<f32> = (1..=9).map(|x| x as f32).collect();
+        let eye: Vec<f32> = vec![
+            1.0, 0.0, 0.0,
+            0.0, 1.0, 0.0,
+            0.0, 0.0, 1.0,
+        ];
+        let result = MuonOptimizer::matmul(&a, &eye, 3, 3, 3);
+        for (r, e) in result.iter().zip(a.iter()) {
+            assert!((r - e).abs() < 1e-6, "A @ I != A at some element");
+        }
+    }
+
+    #[test]
+    fn test_matmul_2x2() {
+        // [[1,2],[3,4]] @ [[5,6],[7,8]] = [[19,22],[43,50]]
+        let a = vec![1.0_f32, 2.0, 3.0, 4.0];
+        let b = vec![5.0_f32, 6.0, 7.0, 8.0];
+        let c = MuonOptimizer::matmul(&a, &b, 2, 2, 2);
+        let expected = vec![19.0_f32, 22.0, 43.0, 50.0];
+        for (r, e) in c.iter().zip(expected.iter()) {
+            assert!((r - e).abs() < 1e-4, "matmul result {r} != expected {e}");
+        }
+    }
+
+    // -------------------------------------------------------------------------
+    // Newton-Schulz
+    // -------------------------------------------------------------------------
+
+    #[test]
+    fn test_newton_schulz_reduces_norm() {
+        // For a random-ish 4×4 matrix the output Frobenius norm should be ≤ 1.1
+        // (it converges toward sqrt(min(m,n)) ≈ 2 for square matrices,
+        //  but after normalization by ||G||_F the iteration starts from norm≈1).
+        let opt = MuonOptimizer::new(MuonConfig::default());
+        let g: Vec<f32> = (1..=16).map(|x| x as f32 * 0.1).collect();
+        let out = opt.newton_schulz5(&g, 4, 4);
+        let norm = MuonOptimizer::frobenius_norm(&out);
+        // The quintic iteration keeps the Frobenius norm bounded; empirically
+        // the result lands well below 4 (= sqrt(4)*sqrt(4)) for this input.
+        assert!(norm < 4.5, "NS output Frobenius norm {norm} unexpectedly large");
+        assert!(norm > 0.0, "NS output is all-zero");
+    }
+
+    #[test]
+    fn test_newton_schulz_1x1() {
+        // For a 1×1 matrix G = [[v]], X_0 = v/|v|.
+        // The quintic polynomial f(x) = a·x + b·x³ + c·x⁵ with the Kosson
+        // coefficients does NOT have x=1 as a fixed point; instead it maps
+        // x0=1.0 → ≈0.701 on the first iteration.  After 5 iterations the
+        // value oscillates in (0.6, 1.2).
+        //
+        // The key correctness properties are:
+        //   * output has the same sign as input
+        //   * output magnitude is bounded in (0.0, 2.0) (not divergent)
+        let opt = MuonOptimizer::new(MuonConfig::default());
+        let g_pos = vec![5.0_f32];
+        let out_pos = opt.newton_schulz5(&g_pos, 1, 1);
+        assert_eq!(out_pos.len(), 1, "output must be length 1");
+        assert!(out_pos[0] > 0.0, "positive input → positive output");
+        assert!(out_pos[0] < 2.0, "1×1 NS output should not diverge (got {})", out_pos[0]);
+
+        // Negative input → negative output (sign preservation)
+        let g_neg = vec![-3.0_f32];
+        let out_neg = opt.newton_schulz5(&g_neg, 1, 1);
+        assert!(out_neg[0] < 0.0, "negative input → negative output");
+    }
+
+    #[test]
+    fn test_ns_output_shape() {
+        let opt = MuonOptimizer::new(MuonConfig::default());
+        let g = vec![1.0_f32; 12]; // 3×4
+        let out = opt.newton_schulz5(&g, 3, 4);
+        assert_eq!(out.len(), 12, "NS output must have the same number of elements as input");
+    }
+
+    // -------------------------------------------------------------------------
+    // Optimizer step
+    // -------------------------------------------------------------------------
+
+    #[test]
+    fn test_step_1d_no_ns() {
+        // For cols==1 the update must equal the Nesterov gradient directly,
+        // not the NS-orthogonalized version.  We verify by checking that the
+        // param changes by exactly lr * nesterov_g (no NS distortion).
+        let mut opt = MuonOptimizer::with_lr(0.1);
+        opt.register("b", 4, 1);
+
+        let params = vec![1.0_f32; 4];
+        let grad = vec![1.0_f32; 4];
+        // First step: m_0=0, so m_1 = 0.95*0 + 1.0 = 1.0
+        //             nesterov = 1.0 + 0.95*1.0 = 1.95
+        //             update = params - lr * nesterov = 1.0 - 0.1 * 1.95 = 0.805
+        let updated = opt.step("b", &params, &grad);
+        for &u in &updated {
+            assert!((u - 0.805_f32).abs() < 1e-5, "1-D step wrong: got {u}, expected 0.805");
+        }
+    }
+
+    #[test]
+    fn test_step_2d_updates_params() {
+        let mut opt = MuonOptimizer::with_lr(0.02);
+        opt.register("W", 4, 4);
+
+        let params = vec![1.0_f32; 16];
+        let grad = vec![0.1_f32; 16];
+        let updated = opt.step("W", &params, &grad);
+
+        assert_eq!(updated.len(), 16, "output must have same length as input");
+        // Params should have changed
+        let changed = params.iter().zip(updated.iter()).any(|(p, u)| (p - u).abs() > 1e-9);
+        assert!(changed, "2-D step should change params");
+    }
+
+    #[test]
+    fn test_momentum_accumulates() {
+        // Verify that the momentum buffer values grow after multiple steps with
+        // the same gradient.  We use a 1-D parameter (cols=1, no NS) so that
+        // the accumulation is directly visible in the update magnitude.
+        //
+        // For cols=1 (plain-SGD path):
+        //   step 1: m1 = 0.95*0 + 1 = 1.0 ;  nesterov = 1 + 0.95*1.0 = 1.95
+        //   step 2: m2 = 0.95*1 + 1 = 1.95 ; nesterov = 1 + 0.95*1.95 = 2.8525
+        // So |update| grows monotonically.
+        let mut opt = MuonOptimizer::with_lr(0.001);
+        opt.register("b", 4, 1); // 1-D: no NS normalization
+
+        let params = vec![10.0_f32; 4];
+        let grad = vec![1.0_f32; 4];
+
+        let p1 = opt.step("b", &params, &grad);
+        let p2 = opt.step("b", &p1, &grad);
+
+        let delta1: f32 = params.iter().zip(p1.iter()).map(|(a, b)| (a - b).abs()).sum();
+        let delta2: f32 = p1.iter().zip(p2.iter()).map(|(a, b)| (a - b).abs()).sum();
+
+        assert!(
+            delta2 > delta1,
+            "momentum should cause step 2 delta ({delta2}) > step 1 delta ({delta1})"
+        );
+    }
+
+    #[test]
+    fn test_step_reduces_loss_direction() {
+        // For a simple quadratic loss L = 0.5 * ||W||², gradient = W,
+        // the update should move W toward zero (loss-decreasing direction).
+        let mut opt = MuonOptimizer::with_lr(0.02);
+        opt.register("W", 3, 3);
+
+        let params = vec![2.0_f32; 9];
+        let grad = params.clone(); // gradient of 0.5||W||² is W
+
+        let updated = opt.step("W", &params, &grad);
+
+        let norm_before: f32 = params.iter().map(|x| x * x).sum::<f32>().sqrt();
+        let norm_after: f32 = updated.iter().map(|x| x * x).sum::<f32>().sqrt();
+
+        assert!(
+            norm_after < norm_before,
+            "update should reduce ||W||: before={norm_before}, after={norm_after}"
+        );
+    }
+
+    #[test]
+    fn test_weight_decay_shrinks_params() {
+        // With non-zero weight decay and zero gradients, params should shrink.
+        let mut opt = MuonOptimizer::new(MuonConfig {
+            lr: 0.01,
+            weight_decay: 0.1,
+            ..MuonConfig::default()
+        });
+        opt.register("W", 3, 3);
+
+        let params = vec![1.0_f32; 9];
+        let grad = vec![0.0_f32; 9];
+
+        let updated = opt.step("W", &params, &grad);
+
+        // Every element should be smaller than 1.0 due to weight decay
+        for &u in &updated {
+            assert!(u < 1.0_f32, "weight decay should shrink params, got {u}");
+        }
+    }
+
+    #[test]
+    #[should_panic(expected = "not registered")]
+    fn test_register_then_step_panics_if_unregistered() {
+        let mut opt = MuonOptimizer::with_lr(0.01);
+        // NOTE: deliberately NOT calling opt.register(...)
+        let params = vec![1.0_f32; 4];
+        let grad = vec![0.1_f32; 4];
+        let _ = opt.step("missing_param", &params, &grad);
+    }
+
+    #[test]
+    fn test_step_zero_grad() {
+        // Zero gradient: nesterov grad after one step = β * m_prev = β * 0 = 0
+        // so params should barely change (change only from weight decay = 0 default).
+        let mut opt = MuonOptimizer::with_lr(0.02);
+        opt.register("W", 3, 3);
+
+        let params = vec![1.0_f32; 9];
+        let grad = vec![0.0_f32; 9];
+
+        let updated = opt.step("W", &params, &grad);
+
+        // With zero grad and no weight decay, params should be unchanged
+        for (p, u) in params.iter().zip(updated.iter()) {
+            assert!((p - u).abs() < 1e-6, "zero grad should not change params: {p} → {u}");
+        }
+    }
+}