spb/forge Public MIT

1// Author: Simon-Pierre Boucher — contact@spboucher.ai2#include "train/optimizer.h"34#include "ops/cpu/cpu_ops.h"5#include "ops/metal/metal_ops.h"6#include "ops/ops.h"78#include <cmath>9#include <unordered_set>1011namespace forge::train {1213namespace {1415// Newton-Schulz quintic coefficients (Keller Jordan's Muon: tuned so the16// iteration's fixed point spreads singular values toward 1 fast; 5 rounds17// suffice at bf16-level accuracy, and we run f32).18constexpr float kNsA = 3.4445f;19constexpr float kNsB = -4.7750f;20constexpr float kNsC = 2.0315f;21constexpr int kNsIters = 5;22constexpr float kNsEps = 1e-7f;2324// Muon updates hidden 2-D matrices only; embeddings and the (possibly tied)25// lm_head keep AdamW, matching the reference MuonWithAuxAdam split.26bool muon_eligible(const std::string& name, const Var& p) {27    if (p.value().ndim() != 2) return false;28    if (name.find("emb") != std::string::npos) return false;     // tok_emb/pos_emb29    if (name.find("lm_head") != std::string::npos) return false;30    return true;31}3233// LR adjustment for rectangular matrices (Muon reference impl).34float muon_adj(const Tensor& w) {35    const float r = float(w.size(0)), c = float(w.size(1));36    return std::sqrt(std::max(1.0f, r / c));37}3839float frobenius(const Tensor& x) {40    double sq = 0.0;41    const float* p = x.data<float>();42    for (int64_t i = 0; i < x.numel(); ++i) sq += double(p[i]) * double(p[i]);43    return float(std::sqrt(sq));44}4546} // namespace4748Optimizer::Optimizer(const std::vector<std::pair<std::string, Var>>& named_params,49                     Options opts)50    : opts_(opts) {51    std::unordered_set<const void*> seen;52    for (const auto& [name, p] : named_params) {53        if (!p.defined() || !p.requires_grad()) continue;54        if (!seen.insert(p.id()).second) continue; // tied param, already tracked55        params_.push_back(p);56        decay_.push_back(p.value().ndim() >= 2);57        muon_.push_back(opts_.kind == "muon" && muon_eligible(name, p));58    }59}6061float Optimizer::clip_global_norm(float max_norm) {62    double sq = 0.0;63    for (const Var& p : params_) {64        if (!p.has_grad()) continue;65        const float* g = p.grad().data<float>();66        for (int64_t i = 0; i < p.grad().numel(); ++i) sq += double(g[i]) * double(g[i]);67    }68    const float norm = float(std::sqrt(sq));69    if (max_norm > 0.0f && norm > max_norm) {70        const float s = max_norm / norm;71        for (const Var& p : params_) {72            if (!p.has_grad()) continue;73            float* g = p.grad().data<float>();74            for (int64_t i = 0; i < p.grad().numel(); ++i) g[i] *= s;75        }76    }77    return norm;78}7980void Optimizer::ensure_state() {81    if (!m_.empty()) return;82    m_.reserve(params_.size());83    v_.reserve(params_.size());84    for (const Var& p : params_) {85        m_.push_back(Tensor::zeros(p.value().shape()));86        v_.push_back(Tensor::zeros(p.value().shape()));87    }88}8990float Optimizer::step_with_clip(float lr, float max_norm) {91    if (ops::backend() == ops::Backend::CPU) {92        const float norm = clip_global_norm(max_norm);93        step(lr);94        return norm;95    }9697    // Metal path. Encode per-tensor sum-of-squares into one partials buffer,98    // then sync — this is also the boundary where the step's loss becomes99    // readable.100    // These are hundreds of independent single-threadgroup reductions (one per101    // parameter tensor), which a serial encoder pointlessly serializes.102    Tensor partials = Tensor::empty({int64_t(params_.size())});103    {104        metal::ConcurrentRegion region;105        for (size_t i = 0; i < params_.size(); ++i) {106            Tensor slot = partials.slice0(int64_t(i), 1);107            metal::sumsq(params_[i].grad(), slot);108        }109    }110    metal::sync();111112    double sq = 0.0;113    for (size_t i = 0; i < params_.size(); ++i) sq += double(partials.data<float>()[i]);114    const float norm = float(std::sqrt(sq));115    const float grad_scale =116        (max_norm > 0.0f && norm > max_norm) ? max_norm / norm : 1.0f;117118    ensure_state(); // CPU-side zeros are safe here: the stream is closed119    ++t_;120    {121        // Each parameter's AdamW update touches only its own w/g/m/v — independent.122        metal::ConcurrentRegion region;123        for (size_t i = 0; i < params_.size(); ++i) {124            if (muon_[i]) continue;125            const float wd = decay_[i] ? opts_.weight_decay : 0.0f;126            metal::adamw_step(params_[i].value(), params_[i].grad(), m_[i], v_[i], lr,127                              opts_.beta1, opts_.beta2, t_, opts_.eps, wd, grad_scale);128        }129    }130131    // Muon phase 1: momentum + nesterov direction + its Frobenius norm.132    // (Serial encoder: each param's chain has read-after-write dependencies;133    // cross-param over-ordering is harmless — these are tiny dispatches.)134    std::vector<Tensor> xs(params_.size());135    Tensor mu_partials;136    std::vector<int64_t> mu_slot(params_.size(), -1);137    int64_t n_muon = 0;138    for (size_t i = 0; i < params_.size(); ++i)139        if (muon_[i]) mu_slot[i] = n_muon++;140    if (n_muon > 0) {141        mu_partials = Tensor::empty({n_muon});142        const float beta = opts_.muon_momentum;143        for (size_t i = 0; i < params_.size(); ++i) {144            if (!muon_[i]) continue;145            const Tensor& g = params_[i].grad();146            Tensor gs = Tensor::empty(g.shape());147            xs[i] = Tensor::empty(g.shape());148            metal::scale(g, grad_scale, gs);      // g̃ = clip-scaled grad149            metal::scale(m_[i], beta, m_[i]);150            metal::accumulate(m_[i], gs);         // m = β·m + g̃151            metal::scale(m_[i], beta, xs[i]);152            metal::accumulate(xs[i], gs);         // X = g̃ + β·m (nesterov)153            Tensor slot = mu_partials.slice0(mu_slot[i], 1);154            metal::sumsq(xs[i], slot);155        }156        metal::sync();157158        // Muon phase 2: normalize, orthogonalize (Newton-Schulz via the159        // existing matmul kernels), apply. For rows > cols the iteration runs160        // on the implicit transpose: A = XᵀX and X ← aX + XB, which is the161        // transpose of the tall-side recurrence — no materialized transpose.162        const float lr_m = lr * opts_.muon_lr_ratio;163        for (size_t i = 0; i < params_.size(); ++i) {164            if (!muon_[i]) continue;165            Tensor& X = xs[i];166            const int64_t r = X.size(0), c = X.size(1);167            const bool tall = r > c;168            const int64_t k = tall ? c : r;169            const float fro = std::sqrt(mu_partials.data<float>()[mu_slot[i]]);170            metal::scale(X, 1.0f / (fro + kNsEps), X);171172            Tensor A = Tensor::empty({k, k});173            Tensor B = Tensor::empty({k, k});174            Tensor X2 = Tensor::empty({r, c});175            for (int it = 0; it < kNsIters; ++it) {176                if (tall) metal::matmul(X, X, A, true, false);   // A = XᵀX177                else      metal::matmul(X, X, A, false, true);   // A = XXᵀ178                metal::matmul(A, A, B);179                metal::scale(B, kNsC, B);180                metal::scale(A, kNsB, A);181                metal::accumulate(B, A);                          // B = bA + cA²182                if (tall) metal::matmul(X, B, X2);                // X2 = XB183                else      metal::matmul(B, X, X2);                // X2 = BX184                metal::scale(X, kNsA, X);185                metal::accumulate(X, X2);                         // X = aX + X2186            }187            metal::scale(params_[i].value(), 1.0f - lr_m * opts_.weight_decay,188                         params_[i].value());189            metal::scale(X, -lr_m * muon_adj(params_[i].value()), X);190            metal::accumulate(params_[i].value(), X);191        }192    }193194    metal::sync(); // weights final before the next step's CPU-side zero_grad195    return norm;196}197198void Optimizer::adamw_cpu(size_t pi, float lr) {199    const float bc1 = 1.0f - std::pow(opts_.beta1, float(t_));200    const float bc2 = 1.0f - std::pow(opts_.beta2, float(t_));201    Var& p = params_[pi];202    const float wd = decay_[pi] ? opts_.weight_decay : 0.0f;203    float* w = p.value().data<float>();204    const float* g = p.grad().data<float>();205    float* m = m_[pi].data<float>();206    float* v = v_[pi].data<float>();207    const int64_t n = p.value().numel();208    for (int64_t i = 0; i < n; ++i) {209        m[i] = opts_.beta1 * m[i] + (1.0f - opts_.beta1) * g[i];210        v[i] = opts_.beta2 * v[i] + (1.0f - opts_.beta2) * g[i] * g[i];211        const float mhat = m[i] / bc1;212        const float vhat = v[i] / bc2;213        w[i] -= lr * (mhat / (std::sqrt(vhat) + opts_.eps) + wd * w[i]);214    }215}216217void Optimizer::muon_cpu(size_t pi, float lr) {218    Var& p = params_[pi];219    const float beta = opts_.muon_momentum;220    const float* g = p.grad().data<float>();221    float* m = m_[pi].data<float>();222    const int64_t n = p.value().numel();223    const int64_t r = p.value().size(0), c = p.value().size(1);224    const bool tall = r > c;225    const int64_t k = tall ? c : r;226227    Tensor X = Tensor::empty({r, c});228    float* x = X.data<float>();229    for (int64_t i = 0; i < n; ++i) {230        m[i] = beta * m[i] + g[i];231        x[i] = g[i] + beta * m[i]; // nesterov232    }233    const float fro = frobenius(X);234    for (int64_t i = 0; i < n; ++i) x[i] /= (fro + kNsEps);235236    Tensor A = Tensor::empty({k, k});237    Tensor B = Tensor::empty({k, k});238    Tensor X2 = Tensor::empty({r, c});239    for (int it = 0; it < kNsIters; ++it) {240        if (tall) cpu::matmul(X, X, A, true, false);241        else      cpu::matmul(X, X, A, false, true);242        cpu::matmul(A, A, B);243        float* pa = A.data<float>();244        float* pb = B.data<float>();245        for (int64_t i = 0; i < k * k; ++i) pb[i] = kNsB * pa[i] + kNsC * pb[i];246        if (tall) cpu::matmul(X, B, X2);247        else      cpu::matmul(B, X, X2);248        const float* p2 = X2.data<float>();249        for (int64_t i = 0; i < n; ++i) x[i] = kNsA * x[i] + p2[i];250    }251252    const float lr_m = lr * opts_.muon_lr_ratio;253    const float adj = muon_adj(p.value());254    float* w = p.value().data<float>();255    for (int64_t i = 0; i < n; ++i)256        w[i] = w[i] * (1.0f - lr_m * opts_.weight_decay) - lr_m * adj * x[i];257}258259void Optimizer::step(float lr) {260    ensure_state();261    ++t_; // 1-based: bias correction divides by (1 - beta^t)262    for (size_t pi = 0; pi < params_.size(); ++pi) {263        if (!params_[pi].has_grad()) continue;264        if (muon_[pi]) muon_cpu(pi, lr);265        else adamw_cpu(pi, lr);266    }267}268269void Optimizer::zero_grad() {270    for (const Var& p : params_) p.zero_grad();271}272273} // namespace forge::train274