spb/forge Public MIT

1// Author: Simon-Pierre Boucher — contact@spboucher.ai2#pragma once34#include "core/autograd.h"56#include <string>7#include <utility>8#include <vector>910namespace forge::train {1112// Two optimizer kinds behind one interface:13//  - "adamw": AdamW with decoupled weight decay on every parameter14//    (llm.c/PyTorch convention: eps outside sqrt, wd only on dim>=2 params).15//  - "muon": Keller Jordan's Muon on 2-D hidden matrices — momentum16//    orthogonalized by 5 Newton-Schulz iterations — with AdamW kept for17//    embeddings / lm_head / 1-D params (the standard Muon split). The Muon18//    group's LR is lr * muon_lr_ratio so both groups follow one schedule.19// Deduplicates tied parameters by Var::id(). m_ doubles as Muon's momentum20// buffer; v_ stays zero for Muon params so the checkpoint format is21// identical across kinds.22class Optimizer {23public:24    struct Options {25        std::string kind = "adamw";  // "adamw" | "muon"26        float beta1 = 0.9f;27        float beta2 = 0.95f;28        float eps = 1e-8f;29        float weight_decay = 0.1f;30        float muon_momentum = 0.95f;31        float muon_lr_ratio = 1.0f;  // muon lr = lr * ratio32    };3334    Optimizer(const std::vector<std::pair<std::string, Var>>& named_params, Options opts);3536    // Global-norm gradient clip: returns the pre-clip norm and scales all37    // grads by min(1, max_norm/norm). No-op when max_norm <= 0. (CPU path.)38    float clip_global_norm(float max_norm);3940    void step(float lr);41    void zero_grad();4243    // Backend-routed step: clip + update in one call, returns the pre-clip44    // grad norm. On Metal it expects backward() already encoded: it encodes45    // per-tensor sumsq, syncs (this is the step's loss-readback boundary),46    // folds the clip factor into the update, and syncs again at the end.47    // Muon adds one extra sync: the Newton-Schulz input must be normalized48    // by its Frobenius norm, which the CPU reads between the two phases.49    float step_with_clip(float lr, float max_norm);5051    // Allocate m/v now (checkpoint loading needs the buffers to exist).52    void ensure_state();5354    int64_t t() const { return t_; }55    // Checkpoint access (M4): moments in parameter order.56    std::vector<Tensor>& m() { return m_; }57    std::vector<Tensor>& v() { return v_; }58    void set_t(int64_t t) { t_ = t; }59    const std::vector<Var>& params() const { return params_; }6061private:62    // AdamW update for one param (CPU path).63    void adamw_cpu(size_t pi, float lr);64    // Muon update for one param (CPU path); grads already clipped.65    void muon_cpu(size_t pi, float lr);6667    Options opts_;68    std::vector<Var> params_;69    std::vector<bool> decay_;   // dim >= 270    std::vector<bool> muon_;    // kind=="muon" && 2-D && not embedding/head71    std::vector<Tensor> m_, v_; // f32, allocated lazily at first step72    int64_t t_ = 0;73};7475// Historical name; checkpoints and tests predate the Muon mode.76using AdamW = Optimizer;7778} // namespace forge::train79