SPB Git

spb/forge Public MIT

Forge — LLM training from scratch in pure C++20 + Metal on Apple Silicon.

llmtrainingcppmetalapple-silicon
C++ 61.2% C 23% Python 7.6% TeX 7.2% CMake 1.1%

Add configurable training modes: Muon optimizer and WSD schedule

- optimizer "adamw" | "muon": Newton-Schulz orthogonalized momentum on 2-D
  hidden matrices (composed from the existing matmul kernels on Metal),
  AdamW kept for embeddings/head/1-D params; muon_lr follows the lr schedule
- schedule "cosine" | "wsd": warmup-stable-decay with 1-sqrt cooldown,
  extendable runs, wsd_decay_frac
- gpt-50m base config + Muon+WSD and MobileLLM-style deep-and-thin variants

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
simon-pierre boucher committed 5 days ago (Aug 5, 2026) parent 75680a9

Showing 8 changed files with +436 and −39

added configs/gpt-50m-deep.json +38 −0
@@ -0,0 +1,38 @@
1 +{
2 + "_comment": "MobileLLM-style deep-and-thin variant of gpt-50m: 20 layers x d_model 448 (aspect ratio ~22 vs 64 for the wide config), head_dim 64. ~50.6M params. Same data/steps/optimizer as gpt-50m so the only variable is the shape.",
3 + "model": {
4 + "name": "gpt-50m-deep",
5 + "n_layers": 20,
6 + "d_model": 448,
7 + "n_heads": 7,
8 + "n_kv_heads": 7,
9 + "d_ff": 1216,
10 + "vocab_size": 4096,
11 + "context_length": 1024,
12 + "tied_embeddings": true,
13 + "use_rope": true,
14 + "rope_theta": 10000.0,
15 + "norm": "rmsnorm",
16 + "norm_eps": 1e-06,
17 + "activation": "swiglu",
18 + "dropout": 0.0
19 + },
20 + "train": {
21 + "lr": 0.0005,
22 + "min_lr_ratio": 0.1,
23 + "warmup_steps": 117,
24 + "max_steps": 1170,
25 + "beta1": 0.9,
26 + "beta2": 0.95,
27 + "eps": 1e-08,
28 + "weight_decay": 0.1,
29 + "grad_clip": 1.0,
30 + "batch_size": 8,
31 + "grad_accum_steps": 8,
32 + "precision": "f32",
33 + "checkpoint_every": 200,
34 + "eval_every": 100,
35 + "eval_batches": 20,
36 + "seed": 1337
37 + }
38 +}
added configs/gpt-50m-muon.json +43 −0
@@ -0,0 +1,43 @@
1 +{
2 + "_comment": "gpt-50m trained with Muon (NS-orthogonalized momentum on hidden matrices, AdamW on embeddings/head) + WSD schedule (flat plateau, 15% 1-sqrt cooldown). muon_lr 0.02 is the NanoGPT-speedrun default; lr applies to the AdamW group. Same data/steps as gpt-50m for A/B comparison.",
3 + "model": {
4 + "name": "gpt-50m-muon",
5 + "n_layers": 10,
6 + "d_model": 640,
7 + "n_heads": 10,
8 + "n_kv_heads": 10,
9 + "d_ff": 1728,
10 + "vocab_size": 4096,
11 + "context_length": 1024,
12 + "tied_embeddings": true,
13 + "use_rope": true,
14 + "rope_theta": 10000.0,
15 + "norm": "rmsnorm",
16 + "norm_eps": 1e-06,
17 + "activation": "swiglu",
18 + "dropout": 0.0
19 + },
20 + "train": {
21 + "lr": 0.0005,
22 + "min_lr_ratio": 0.1,
23 + "warmup_steps": 117,
24 + "max_steps": 1170,
25 + "schedule": "wsd",
26 + "wsd_decay_frac": 0.15,
27 + "optimizer": "muon",
28 + "muon_lr": 0.02,
29 + "muon_momentum": 0.95,
30 + "beta1": 0.9,
31 + "beta2": 0.95,
32 + "eps": 1e-08,
33 + "weight_decay": 0.1,
34 + "grad_clip": 1.0,
35 + "batch_size": 8,
36 + "grad_accum_steps": 8,
37 + "precision": "f32",
38 + "checkpoint_every": 200,
39 + "eval_every": 100,
40 + "eval_batches": 20,
41 + "seed": 1337
42 + }
43 +}
added configs/gpt-50m.json +38 −0
@@ -0,0 +1,38 @@
1 +{
2 + "_comment": "~52M params, sized between gpt-25m and gpt-100m. vocab 4096 to match the local tok4096 train.bin (19.14M tokens). batch_size is the MICRO-batch: 8 x 8 x 1024 = 65536 tokens/step; 1170 steps = ~4 epochs over the 19.14M-token set. Warmup is 10% of the run. precision is parsed but not yet honored - all kernels are f32.",
3 + "model": {
4 + "name": "gpt-50m",
5 + "n_layers": 10,
6 + "d_model": 640,
7 + "n_heads": 10,
8 + "n_kv_heads": 10,
9 + "d_ff": 1728,
10 + "vocab_size": 4096,
11 + "context_length": 1024,
12 + "tied_embeddings": true,
13 + "use_rope": true,
14 + "rope_theta": 10000.0,
15 + "norm": "rmsnorm",
16 + "norm_eps": 1e-06,
17 + "activation": "swiglu",
18 + "dropout": 0.0
19 + },
20 + "train": {
21 + "lr": 0.0005,
22 + "min_lr_ratio": 0.1,
23 + "warmup_steps": 117,
24 + "max_steps": 1170,
25 + "beta1": 0.9,
26 + "beta2": 0.95,
27 + "eps": 1e-08,
28 + "weight_decay": 0.1,
29 + "grad_clip": 1.0,
30 + "batch_size": 8,
31 + "grad_accum_steps": 8,
32 + "precision": "f32",
33 + "checkpoint_every": 200,
34 + "eval_every": 100,
35 + "eval_batches": 20,
36 + "seed": 1337
37 + }
38 +}
modified src/nn/config.h +68 −2
@@ -28,6 +28,21 @@ struct ModelConfig {
28 28 float norm_eps = 1e-6f;
29 29 std::string activation = "swiglu"; // "swiglu" | "gelu"
30 30 float dropout = 0.0f;
31 + // Architecture-variant knobs (all off by default = the classic recipe):
32 + bool qk_norm = false; // RMSNorm on q/k per head before RoPE (Qwen3/Gemma3)
33 + float final_softcap = 0.0f; // logits = cap*tanh(logits/cap) (Gemma2: 30)
34 + bool scale_embeddings = false;// x *= sqrt(d_model) after embedding (Gemma)
35 + // Quantization-aware training: linear weights are fake-quantized each
36 + // forward (per-row scales, straight-through estimator in backward).
37 + // Master weights and the optimizer stay f32.
38 + std::string quant = "none"; // "none" | "int8" | "ternary"
39 + // Mixture-of-experts MLP: n_experts > 0 replaces each block's MLP with
40 + // `n_experts` expert MLPs (d_ff each) + a top-k softmax router. v1
41 + // computes every expert densely and masks — correctness first.
42 + int64_t n_experts = 0; // 0 => dense MLP
43 + int64_t moe_top_k = 2;
44 + float moe_aux_weight = 0.01f; // load-balance loss: E * sum_e mean_gate_e^2
45 + int64_t n_shared_experts = 0; // always-active experts (DeepSeek/Kimi style)
31 46
32 47 int64_t head_dim() const { return d_model / n_heads; }
33 48
@@ -37,10 +52,14 @@ struct ModelConfig {
37 52 const int64_t attn = d_model * d_model // wq
38 53 + 2 * d_model * n_kv_heads * hd // wk, wv
39 54 + d_model * d_model; // wo
40 − const int64_t mlp = (activation == "swiglu")
55 + const int64_t mlp_one = (activation == "swiglu")
41 56 ? 3 * d_model * d_ff
42 57 : 2 * d_model * d_ff;
43 − const int64_t norms = (norm == "layernorm" ? 2 : 1) * d_model * (2 * n_layers + 1);
58 + const int64_t mlp = n_experts > 0
59 + ? (n_experts + n_shared_experts) * mlp_one + n_experts * d_model
60 + : mlp_one; // experts + router
61 + int64_t norms = (norm == "layernorm" ? 2 : 1) * d_model * (2 * n_layers + 1);
62 + if (qk_norm) norms += 2 * head_dim() * n_layers;
44 63 int64_t total = n_layers * (attn + mlp) + norms + vocab_size * d_model;
45 64 if (!tied_embeddings) total += vocab_size * d_model;
46 65 if (!use_rope) total += context_length * d_model;
@@ -53,6 +72,14 @@ struct TrainConfig {
53 72 float min_lr_ratio = 0.1f; // min_lr = lr * ratio
54 73 int64_t warmup_steps = 2000;
55 74 int64_t max_steps = 100000;
75 + std::string schedule = "cosine"; // "cosine" | "wsd" (warmup-stable-decay)
76 + float wsd_decay_frac = 0.15f; // WSD: fraction of max_steps in the cooldown
77 + // "muon": Newton-Schulz orthogonalized momentum on 2-D hidden matrices,
78 + // AdamW on embeddings/head/norms (Keller Jordan's Muon). muon_lr is the
79 + // peak LR of the Muon group; it follows the same schedule shape as `lr`.
80 + std::string optimizer = "adamw"; // "adamw" | "muon"
81 + float muon_lr = 0.02f;
82 + float muon_momentum = 0.95f;
56 83 float beta1 = 0.9f;
57 84 float beta2 = 0.95f;
58 85 float eps = 1e-8f;
@@ -62,6 +89,11 @@ struct TrainConfig {
62 89 int64_t grad_accum_steps = 1;
63 90 std::string precision = "f32"; // "f32" | "f16" | "bf16" (compute dtype)
64 91 int64_t checkpoint_every = 1000;
92 + // Also commit each checkpoint's weights into <out>/model.forge — the
93 + // git-style zero-copy format (src/core/fmodel.h). The .bin checkpoint
94 + // stays the resume source of truth (it carries optimizer state).
95 + bool forge_save = true;
96 + std::string forge_dtype = "f32"; // "f32" | "f16" | "bf16"
65 97 int64_t eval_every = 500;
66 98 int64_t eval_batches = 20;
67 99 uint64_t seed = 1337;
@@ -89,6 +121,14 @@ inline void from_json(const nlohmann::json& j, ModelConfig& c) {
89 121 c.norm_eps = j.value("norm_eps", c.norm_eps);
90 122 c.activation = j.value("activation", c.activation);
91 123 c.dropout = j.value("dropout", c.dropout);
124 + c.quant = j.value("quant", c.quant);
125 + c.n_experts = j.value("n_experts", c.n_experts);
126 + c.moe_top_k = j.value("moe_top_k", c.moe_top_k);
127 + c.moe_aux_weight = j.value("moe_aux_weight", c.moe_aux_weight);
128 + c.n_shared_experts = j.value("n_shared_experts", c.n_shared_experts);
129 + c.qk_norm = j.value("qk_norm", c.qk_norm);
130 + c.final_softcap = j.value("final_softcap", c.final_softcap);
131 + c.scale_embeddings = j.value("scale_embeddings", c.scale_embeddings);
92 132
93 133 if (c.d_model % c.n_heads != 0)
94 134 throw std::runtime_error("config: d_model must be divisible by n_heads");
@@ -98,6 +138,16 @@ inline void from_json(const nlohmann::json& j, ModelConfig& c) {
98 138 throw std::runtime_error("config: norm must be rmsnorm or layernorm");
99 139 if (c.activation != "swiglu" && c.activation != "gelu")
100 140 throw std::runtime_error("config: activation must be swiglu or gelu");
141 + if (c.quant != "none" && c.quant != "int8" && c.quant != "ternary")
142 + throw std::runtime_error("config: quant must be none, int8 or ternary");
143 + if (c.n_experts < 0)
144 + throw std::runtime_error("config: n_experts must be >= 0");
145 + if (c.n_experts > 0 && (c.moe_top_k < 1 || c.moe_top_k > c.n_experts))
146 + throw std::runtime_error("config: moe_top_k must be in [1, n_experts]");
147 + if (c.n_shared_experts < 0 || (c.n_shared_experts > 0 && c.n_experts == 0))
148 + throw std::runtime_error("config: n_shared_experts requires n_experts > 0");
149 + if (c.final_softcap < 0.0f)
150 + throw std::runtime_error("config: final_softcap must be >= 0");
101 151 }
102 152
103 153 inline void from_json(const nlohmann::json& j, TrainConfig& c) {
@@ -105,6 +155,11 @@ inline void from_json(const nlohmann::json& j, TrainConfig& c) {
105 155 c.min_lr_ratio = j.value("min_lr_ratio", c.min_lr_ratio);
106 156 c.warmup_steps = j.value("warmup_steps", c.warmup_steps);
107 157 c.max_steps = j.value("max_steps", c.max_steps);
158 + c.schedule = j.value("schedule", c.schedule);
159 + c.wsd_decay_frac = j.value("wsd_decay_frac", c.wsd_decay_frac);
160 + c.optimizer = j.value("optimizer", c.optimizer);
161 + c.muon_lr = j.value("muon_lr", c.muon_lr);
162 + c.muon_momentum = j.value("muon_momentum", c.muon_momentum);
108 163 c.beta1 = j.value("beta1", c.beta1);
109 164 c.beta2 = j.value("beta2", c.beta2);
110 165 c.eps = j.value("eps", c.eps);
@@ -114,10 +169,21 @@ inline void from_json(const nlohmann::json& j, TrainConfig& c) {
114 169 c.grad_accum_steps = j.value("grad_accum_steps", c.grad_accum_steps);
115 170 c.precision = j.value("precision", c.precision);
116 171 c.checkpoint_every = j.value("checkpoint_every", c.checkpoint_every);
172 + c.forge_save = j.value("forge_save", c.forge_save);
173 + c.forge_dtype = j.value("forge_dtype", c.forge_dtype);
117 174 c.eval_every = j.value("eval_every", c.eval_every);
118 175 c.eval_batches = j.value("eval_batches", c.eval_batches);
119 176 c.seed = j.value("seed", c.seed);
120 177 c.deterministic = j.value("deterministic", c.deterministic);
178 +
179 + if (c.schedule != "cosine" && c.schedule != "wsd")
180 + throw std::runtime_error("config: schedule must be cosine or wsd");
181 + if (c.wsd_decay_frac <= 0.0f || c.wsd_decay_frac >= 1.0f)
182 + throw std::runtime_error("config: wsd_decay_frac must be in (0, 1)");
183 + if (c.optimizer != "adamw" && c.optimizer != "muon")
184 + throw std::runtime_error("config: optimizer must be adamw or muon");
185 + if (c.forge_dtype != "f32" && c.forge_dtype != "f16" && c.forge_dtype != "bf16")
186 + throw std::runtime_error("config: forge_dtype must be f32, f16 or bf16");
121 187 }
122 188
123 189 inline Config load_config(const std::string& path) {
modified src/train/optimizer.cpp +173 −24
@@ -1,6 +1,7 @@
1 1 // Author: Simon-Pierre Boucher — contact@spboucher.ai
2 2 #include "train/optimizer.h"
3 3
4 +#include "ops/cpu/cpu_ops.h"
4 5 #include "ops/metal/metal_ops.h"
5 6 #include "ops/ops.h"
6 7
@@ -9,7 +10,43 @@
9 10
10 11 namespace forge::train {
11 12
12 −AdamW::AdamW(const std::vector<std::pair<std::string, Var>>& named_params, Options opts)
13 +namespace {
14 +
15 +// Newton-Schulz quintic coefficients (Keller Jordan's Muon: tuned so the
16 +// iteration's fixed point spreads singular values toward 1 fast; 5 rounds
17 +// suffice at bf16-level accuracy, and we run f32).
18 +constexpr float kNsA = 3.4445f;
19 +constexpr float kNsB = -4.7750f;
20 +constexpr float kNsC = 2.0315f;
21 +constexpr int kNsIters = 5;
22 +constexpr float kNsEps = 1e-7f;
23 +
24 +// Muon updates hidden 2-D matrices only; embeddings and the (possibly tied)
25 +// lm_head keep AdamW, matching the reference MuonWithAuxAdam split.
26 +bool 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_emb
29 + if (name.find("lm_head") != std::string::npos) return false;
30 + return true;
31 +}
32 +
33 +// LR adjustment for rectangular matrices (Muon reference impl).
34 +float 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 +}
38 +
39 +float 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 +}
45 +
46 +} // namespace
47 +
48 +Optimizer::Optimizer(const std::vector<std::pair<std::string, Var>>& named_params,
49 + Options opts)
13 50 : opts_(opts) {
14 51 std::unordered_set<const void*> seen;
15 52 for (const auto& [name, p] : named_params) {
@@ -17,10 +54,11 @@ AdamW::AdamW(const std::vector<std::pair<std::string, Var>>& named_params, Optio
17 54 if (!seen.insert(p.id()).second) continue; // tied param, already tracked
18 55 params_.push_back(p);
19 56 decay_.push_back(p.value().ndim() >= 2);
57 + muon_.push_back(opts_.kind == "muon" && muon_eligible(name, p));
20 58 }
21 59 }
22 60
23 −float AdamW::clip_global_norm(float max_norm) {
61 +float Optimizer::clip_global_norm(float max_norm) {
24 62 double sq = 0.0;
25 63 for (const Var& p : params_) {
26 64 if (!p.has_grad()) continue;
@@ -39,7 +77,7 @@ float AdamW::clip_global_norm(float max_norm) {
39 77 return norm;
40 78 }
41 79
42 −void AdamW::ensure_state() {
80 +void Optimizer::ensure_state() {
43 81 if (!m_.empty()) return;
44 82 m_.reserve(params_.size());
45 83 v_.reserve(params_.size());
@@ -49,7 +87,7 @@ void AdamW::ensure_state() {
49 87 }
50 88 }
51 89
52 −float AdamW::step_with_clip(float lr, float max_norm) {
90 +float Optimizer::step_with_clip(float lr, float max_norm) {
53 91 if (ops::backend() == ops::Backend::CPU) {
54 92 const float norm = clip_global_norm(max_norm);
55 93 step(lr);
@@ -80,44 +118,155 @@ float AdamW::step_with_clip(float lr, float max_norm) {
80 118 ensure_state(); // CPU-side zeros are safe here: the stream is closed
81 119 ++t_;
82 120 {
83 − // Each parameter's update touches only its own w/g/m/v — independent.
121 + // Each parameter's AdamW update touches only its own w/g/m/v — independent.
84 122 metal::ConcurrentRegion region;
85 123 for (size_t i = 0; i < params_.size(); ++i) {
124 + if (muon_[i]) continue;
86 125 const float wd = decay_[i] ? opts_.weight_decay : 0.0f;
87 126 metal::adamw_step(params_[i].value(), params_[i].grad(), m_[i], v_[i], lr,
88 127 opts_.beta1, opts_.beta2, t_, opts_.eps, wd, grad_scale);
89 128 }
90 129 }
130 +
131 + // 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 grad
149 + 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();
157 +
158 + // Muon phase 2: normalize, orthogonalize (Newton-Schulz via the
159 + // existing matmul kernels), apply. For rows > cols the iteration runs
160 + // on the implicit transpose: A = XᵀX and X ← aX + XB, which is the
161 + // 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);
171 +
172 + 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ᵀX
177 + 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 = XB
183 + else metal::matmul(B, X, X2); // X2 = BX
184 + metal::scale(X, kNsA, X);
185 + metal::accumulate(X, X2); // X = aX + X2
186 + }
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 + }
193 +
91 194 metal::sync(); // weights final before the next step's CPU-side zero_grad
92 195 return norm;
93 196 }
94 197
95 −void AdamW::step(float lr) {
96 − ensure_state();
97 − ++t_; // 1-based: bias correction divides by (1 - beta^t)
198 +void Optimizer::adamw_cpu(size_t pi, float lr) {
98 199 const float bc1 = 1.0f - std::pow(opts_.beta1, float(t_));
99 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 +}
100 216
217 +void 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;
226 +
227 + 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]; // nesterov
232 + }
233 + const float fro = frobenius(X);
234 + for (int64_t i = 0; i < n; ++i) x[i] /= (fro + kNsEps);
235 +
236 + 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 + }
251 +
252 + 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 +}
258 +
259 +void Optimizer::step(float lr) {
260 + ensure_state();
261 + ++t_; // 1-based: bias correction divides by (1 - beta^t)
101 262 for (size_t pi = 0; pi < params_.size(); ++pi) {
102 − Var& p = params_[pi];
103 − if (!p.has_grad()) continue;
104 − const float wd = decay_[pi] ? opts_.weight_decay : 0.0f;
105 − float* w = p.value().data<float>();
106 − const float* g = p.grad().data<float>();
107 − float* m = m_[pi].data<float>();
108 − float* v = v_[pi].data<float>();
109 − const int64_t n = p.value().numel();
110 − for (int64_t i = 0; i < n; ++i) {
111 − m[i] = opts_.beta1 * m[i] + (1.0f - opts_.beta1) * g[i];
112 − v[i] = opts_.beta2 * v[i] + (1.0f - opts_.beta2) * g[i] * g[i];
113 − const float mhat = m[i] / bc1;
114 − const float vhat = v[i] / bc2;
115 − w[i] -= lr * (mhat / (std::sqrt(vhat) + opts_.eps) + wd * w[i]);
116 − }
263 + if (!params_[pi].has_grad()) continue;
264 + if (muon_[pi]) muon_cpu(pi, lr);
265 + else adamw_cpu(pi, lr);
117 266 }
118 267 }
119 268
120 −void AdamW::zero_grad() {
269 +void Optimizer::zero_grad() {
121 270 for (const Var& p : params_) p.zero_grad();
122 271 }
123 272
modified src/train/optimizer.h +28 −8
@@ -9,19 +9,29 @@
9 9
10 10 namespace forge::train {
11 11
12 −// AdamW with decoupled weight decay (llm.c/PyTorch convention: eps outside
13 −// sqrt, wd only on dim>=2 params, decay folded into the same update).
14 −// Deduplicates tied parameters by Var::id().
15 −class AdamW {
12 +// Two optimizer kinds behind one interface:
13 +// - "adamw": AdamW with decoupled weight decay on every parameter
14 +// (llm.c/PyTorch convention: eps outside sqrt, wd only on dim>=2 params).
15 +// - "muon": Keller Jordan's Muon on 2-D hidden matrices — momentum
16 +// orthogonalized by 5 Newton-Schulz iterations — with AdamW kept for
17 +// embeddings / lm_head / 1-D params (the standard Muon split). The Muon
18 +// 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 momentum
20 +// buffer; v_ stays zero for Muon params so the checkpoint format is
21 +// identical across kinds.
22 +class Optimizer {
16 23 public:
17 24 struct Options {
25 + std::string kind = "adamw"; // "adamw" | "muon"
18 26 float beta1 = 0.9f;
19 27 float beta2 = 0.95f;
20 28 float eps = 1e-8f;
21 29 float weight_decay = 0.1f;
30 + float muon_momentum = 0.95f;
31 + float muon_lr_ratio = 1.0f; // muon lr = lr * ratio
22 32 };
23 33
24 − AdamW(const std::vector<std::pair<std::string, Var>>& named_params, Options opts);
34 + Optimizer(const std::vector<std::pair<std::string, Var>>& named_params, Options opts);
25 35
26 36 // Global-norm gradient clip: returns the pre-clip norm and scales all
27 37 // grads by min(1, max_norm/norm). No-op when max_norm <= 0. (CPU path.)
@@ -30,11 +40,12 @@ public:
30 40 void step(float lr);
31 41 void zero_grad();
32 42
33 − // Backend-routed step: clip + AdamW in one call, returns the pre-clip
43 + // Backend-routed step: clip + update in one call, returns the pre-clip
34 44 // grad norm. On Metal it expects backward() already encoded: it encodes
35 45 // per-tensor sumsq, syncs (this is the step's loss-readback boundary),
36 − // folds the clip factor into the fused adamw kernel's grad_scale, and
37 − // syncs again after the update.
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 normalized
48 + // by its Frobenius norm, which the CPU reads between the two phases.
38 49 float step_with_clip(float lr, float max_norm);
39 50
40 51 // Allocate m/v now (checkpoint loading needs the buffers to exist).
@@ -48,11 +59,20 @@ public:
48 59 const std::vector<Var>& params() const { return params_; }
49 60
50 61 private:
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);
66 +
51 67 Options opts_;
52 68 std::vector<Var> params_;
53 69 std::vector<bool> decay_; // dim >= 2
70 + std::vector<bool> muon_; // kind=="muon" && 2-D && not embedding/head
54 71 std::vector<Tensor> m_, v_; // f32, allocated lazily at first step
55 72 int64_t t_ = 0;
56 73 };
57 74
75 +// Historical name; checkpoints and tests predate the Muon mode.
76 +using AdamW = Optimizer;
77 +
58 78 } // namespace forge::train
modified src/train/scheduler.h +15 −0
@@ -17,4 +17,19 @@ inline float lr_at(int64_t step, float max_lr, float min_lr, int64_t warmup_step
17 17 return min_lr + coeff * (max_lr - min_lr);
18 18 }
19 19
20 +// Warmup–Stable–Decay (MiniCPM): linear warmup, flat plateau at max_lr, then a
21 +// short cooldown over the final decay_steps. Any plateau checkpoint can be
22 +// resumed and the run extended without re-deciding total_steps up front. The
23 +// cooldown uses the 1-sqrt shape, which beats linear/exponential cooldowns in
24 +// the WSD ablations (arXiv 2404.06395 follow-ups).
25 +inline float lr_wsd(int64_t step, float max_lr, float min_lr, int64_t warmup_steps,
26 + int64_t total_steps, int64_t decay_steps) {
27 + if (step < warmup_steps) return max_lr * float(step + 1) / float(warmup_steps + 1);
28 + const int64_t decay_start = total_steps - decay_steps;
29 + if (step < decay_start) return max_lr;
30 + if (step >= total_steps) return min_lr;
31 + const float ratio = float(step - decay_start) / float(decay_steps);
32 + return min_lr + (max_lr - min_lr) * (1.0f - std::sqrt(ratio));
33 +}
34 +
20 35 } // namespace forge::train
modified src/train/trainer.cpp +33 −5
@@ -1,6 +1,7 @@
1 1 // Author: Simon-Pierre Boucher — contact@spboucher.ai
2 2 #include "train/trainer.h"
3 3
4 +#include "core/fmodel.h"
4 5 #include "ops/metal/metal_ops.h"
5 6 #include "ops/ops.h"
6 7 #include "train/checkpoint.h"
@@ -8,6 +9,7 @@
8 9
9 10 #include <Foundation/Foundation.hpp>
10 11
12 +#include <algorithm>
11 13 #include <chrono>
12 14 #include <cstdio>
13 15 #include <filesystem>
@@ -20,12 +22,15 @@ Trainer::Trainer(Config cfg, const std::string& data_dir, const std::string& out
20 22 std::filesystem::create_directories(out_dir_);
21 23 model_ = std::make_unique<nn::Transformer>(cfg_.model, cfg_.train.seed);
22 24
23 − AdamW::Options opts;
25 + Optimizer::Options opts;
26 + opts.kind = cfg_.train.optimizer;
24 27 opts.beta1 = cfg_.train.beta1;
25 28 opts.beta2 = cfg_.train.beta2;
26 29 opts.eps = cfg_.train.eps;
27 30 opts.weight_decay = cfg_.train.weight_decay;
28 − opt_ = std::make_unique<AdamW>(model_->named_parameters(), opts);
31 + opts.muon_momentum = cfg_.train.muon_momentum;
32 + opts.muon_lr_ratio = cfg_.train.muon_lr / cfg_.train.lr;
33 + opt_ = std::make_unique<Optimizer>(model_->named_parameters(), opts);
29 34
30 35 train_data_ = std::make_unique<DataLoader>(data_dir + "/train.bin",
31 36 cfg_.model.context_length,
@@ -64,6 +69,22 @@ void Trainer::save(int64_t step) {
64 69 save_checkpoint(out_dir_ + "/ckpt_latest.bin", model_->named_parameters(), opt_.get(),
65 70 meta);
66 71 std::printf("checkpoint saved: %s\n", (out_dir_ + name).c_str());
72 +
73 + // Native .forge commit: the run's whole weight history lives in one
74 + // git-style repo; content addressing means an unchanged tensor (frozen,
75 + // masked, tied) is never rewritten.
76 + if (cfg_.train.forge_save) {
77 + fmodel::SaveOptions fo;
78 + fo.dtype = cfg_.train.forge_dtype == "f16" ? DType::F16
79 + : cfg_.train.forge_dtype == "bf16" ? DType::BF16
80 + : DType::F32;
81 + char tag[32];
82 + std::snprintf(tag, sizeof(tag), "step-%06lld", static_cast<long long>(step));
83 + fo.tag = tag;
84 + fo.step = step;
85 + fmodel::save(out_dir_ + "/model.forge", config_json_,
86 + model_->named_parameters(), fo);
87 + }
67 88 }
68 89
69 90 void Trainer::train(const std::string& resume_from) {
@@ -87,18 +108,25 @@ void Trainer::train(const std::string& resume_from) {
87 108 const int64_t min_lr_steps = tc.max_steps;
88 109 const float min_lr = tc.lr * tc.min_lr_ratio;
89 110 const int64_t tokens_per_step = B * T * tc.grad_accum_steps;
111 + const bool wsd = tc.schedule == "wsd";
112 + const int64_t wsd_decay_steps =
113 + std::max<int64_t>(1, int64_t(float(tc.max_steps) * tc.wsd_decay_frac));
90 114
91 − std::printf("training %s: %lld params, %lld steps, %lld tokens/step, backend=%s\n",
115 + std::printf("training %s: %lld params, %lld steps, %lld tokens/step, backend=%s, "
116 + "opt=%s, sched=%s\n",
92 117 cfg_.model.name.c_str(),
93 118 static_cast<long long>(cfg_.model.num_params()),
94 119 static_cast<long long>(tc.max_steps),
95 120 static_cast<long long>(tokens_per_step),
96 − ops::backend() == ops::Backend::Metal ? "metal" : "cpu");
121 + ops::backend() == ops::Backend::Metal ? "metal" : "cpu",
122 + tc.optimizer.c_str(), tc.schedule.c_str());
97 123
98 124 for (int64_t step = start_step; step < tc.max_steps; ++step) {
99 125 NS::AutoreleasePool* pool = NS::AutoreleasePool::alloc()->init();
100 126 const auto t0 = std::chrono::steady_clock::now();
101 − const float lr = lr_at(step, tc.lr, min_lr, tc.warmup_steps, min_lr_steps);
127 + const float lr = wsd
128 + ? lr_wsd(step, tc.lr, min_lr, tc.warmup_steps, tc.max_steps, wsd_decay_steps)
129 + : lr_at(step, tc.lr, min_lr, tc.warmup_steps, min_lr_steps);
102 130
103 131 opt_->zero_grad();
104 132 const bool on_gpu = ops::backend() == ops::Backend::Metal;
105 133