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 3 mo ago (Aug 5, 2026) parent 75680a9 Browse files

8 changed files +436 −39

added configs/gpt-50m-deep.json +38 −0

++{
++  "_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.",
++  "model": {
++    "name": "gpt-50m-deep",
++    "n_layers": 20,
++    "d_model": 448,
++    "n_heads": 7,
++    "n_kv_heads": 7,
++    "d_ff": 1216,
++    "vocab_size": 4096,
++    "context_length": 1024,
++    "tied_embeddings": true,
++    "use_rope": true,
++    "rope_theta": 10000.0,
++    "norm": "rmsnorm",
++    "norm_eps": 1e-06,
++    "activation": "swiglu",
++    "dropout": 0.0
++  },
++  "train": {
++    "lr": 0.0005,
++    "min_lr_ratio": 0.1,
++    "warmup_steps": 117,
++    "max_steps": 1170,
++    "beta1": 0.9,
++    "beta2": 0.95,
++    "eps": 1e-08,
++    "weight_decay": 0.1,
++    "grad_clip": 1.0,
++    "batch_size": 8,
++    "grad_accum_steps": 8,
++    "precision": "f32",
++    "checkpoint_every": 200,
++    "eval_every": 100,
++    "eval_batches": 20,
++    "seed": 1337
++  }
++}

added configs/gpt-50m-muon.json +43 −0

++{
++  "_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.",
++  "model": {
++    "name": "gpt-50m-muon",
++    "n_layers": 10,
++    "d_model": 640,
++    "n_heads": 10,
++    "n_kv_heads": 10,
++    "d_ff": 1728,
++    "vocab_size": 4096,
++    "context_length": 1024,
++    "tied_embeddings": true,
++    "use_rope": true,
++    "rope_theta": 10000.0,
++    "norm": "rmsnorm",
++    "norm_eps": 1e-06,
++    "activation": "swiglu",
++    "dropout": 0.0
++  },
++  "train": {
++    "lr": 0.0005,
++    "min_lr_ratio": 0.1,
++    "warmup_steps": 117,
++    "max_steps": 1170,
++    "schedule": "wsd",
++    "wsd_decay_frac": 0.15,
++    "optimizer": "muon",
++    "muon_lr": 0.02,
++    "muon_momentum": 0.95,
++    "beta1": 0.9,
++    "beta2": 0.95,
++    "eps": 1e-08,
++    "weight_decay": 0.1,
++    "grad_clip": 1.0,
++    "batch_size": 8,
++    "grad_accum_steps": 8,
++    "precision": "f32",
++    "checkpoint_every": 200,
++    "eval_every": 100,
++    "eval_batches": 20,
++    "seed": 1337
++  }
++}

added configs/gpt-50m.json +38 −0

++{
++  "_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.",
++  "model": {
++    "name": "gpt-50m",
++    "n_layers": 10,
++    "d_model": 640,
++    "n_heads": 10,
++    "n_kv_heads": 10,
++    "d_ff": 1728,
++    "vocab_size": 4096,
++    "context_length": 1024,
++    "tied_embeddings": true,
++    "use_rope": true,
++    "rope_theta": 10000.0,
++    "norm": "rmsnorm",
++    "norm_eps": 1e-06,
++    "activation": "swiglu",
++    "dropout": 0.0
++  },
++  "train": {
++    "lr": 0.0005,
++    "min_lr_ratio": 0.1,
++    "warmup_steps": 117,
++    "max_steps": 1170,
++    "beta1": 0.9,
++    "beta2": 0.95,
++    "eps": 1e-08,
++    "weight_decay": 0.1,
++    "grad_clip": 1.0,
++    "batch_size": 8,
++    "grad_accum_steps": 8,
++    "precision": "f32",
++    "checkpoint_every": 200,
++    "eval_every": 100,
++    "eval_batches": 20,
++    "seed": 1337
++  }
++}

modified src/nn/config.h +68 −2

				`@@ -28,6 +28,21 @@ struct ModelConfig {`
28	float norm_eps = 1e-6f;	28	float norm_eps = 1e-6f;
29	std::string activation = "swiglu"; // "swiglu" \| "gelu"	29	std::string activation = "swiglu"; // "swiglu" \| "gelu"
30	float dropout = 0.0f;	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	int64_t head_dim() const { return d_model / n_heads; }	47	int64_t head_dim() const { return d_model / n_heads; }
33		48
				`@@ -37,10 +52,14 @@ struct ModelConfig {`
37	const int64_t attn = d_model * d_model // wq	52	const int64_t attn = d_model * d_model // wq
38	+ 2 * d_model * n_kv_heads * hd // wk, wv	53	+ 2 * d_model * n_kv_heads * hd // wk, wv
39	+ d_model * d_model; // wo	54	+ d_model * d_model; // wo
40	− const int64_t mlp = (activation == "swiglu")	55	+ const int64_t mlp_one = (activation == "swiglu")
41	? 3 * d_model * d_ff	56	? 3 * d_model * d_ff
42	: 2 * d_model * d_ff;	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	int64_t total = n_layers * (attn + mlp) + norms + vocab_size * d_model;	63	int64_t total = n_layers * (attn + mlp) + norms + vocab_size * d_model;
45	if (!tied_embeddings) total += vocab_size * d_model;	64	if (!tied_embeddings) total += vocab_size * d_model;
46	if (!use_rope) total += context_length * d_model;	65	if (!use_rope) total += context_length * d_model;
				`@@ -53,6 +72,14 @@ struct TrainConfig {`
53	float min_lr_ratio = 0.1f; // min_lr = lr * ratio	72	float min_lr_ratio = 0.1f; // min_lr = lr * ratio
54	int64_t warmup_steps = 2000;	73	int64_t warmup_steps = 2000;
55	int64_t max_steps = 100000;	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	float beta1 = 0.9f;	83	float beta1 = 0.9f;
57	float beta2 = 0.95f;	84	float beta2 = 0.95f;
58	float eps = 1e-8f;	85	float eps = 1e-8f;
				`@@ -62,6 +89,11 @@ struct TrainConfig {`
62	int64_t grad_accum_steps = 1;	89	int64_t grad_accum_steps = 1;
63	std::string precision = "f32"; // "f32" \| "f16" \| "bf16" (compute dtype)	90	std::string precision = "f32"; // "f32" \| "f16" \| "bf16" (compute dtype)
64	int64_t checkpoint_every = 1000;	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	int64_t eval_every = 500;	97	int64_t eval_every = 500;
66	int64_t eval_batches = 20;	98	int64_t eval_batches = 20;
67	uint64_t seed = 1337;	99	uint64_t seed = 1337;
				`@@ -89,6 +121,14 @@ inline void from_json(const nlohmann::json& j, ModelConfig& c) {`
89	c.norm_eps = j.value("norm_eps", c.norm_eps);	121	c.norm_eps = j.value("norm_eps", c.norm_eps);
90	c.activation = j.value("activation", c.activation);	122	c.activation = j.value("activation", c.activation);
91	c.dropout = j.value("dropout", c.dropout);	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	if (c.d_model % c.n_heads != 0)	133	if (c.d_model % c.n_heads != 0)
94	throw std::runtime_error("config: d_model must be divisible by n_heads");	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	throw std::runtime_error("config: norm must be rmsnorm or layernorm");	138	throw std::runtime_error("config: norm must be rmsnorm or layernorm");
99	if (c.activation != "swiglu" && c.activation != "gelu")	139	if (c.activation != "swiglu" && c.activation != "gelu")
100	throw std::runtime_error("config: activation must be swiglu or gelu");	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	inline void from_json(const nlohmann::json& j, TrainConfig& c) {	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	c.min_lr_ratio = j.value("min_lr_ratio", c.min_lr_ratio);	155	c.min_lr_ratio = j.value("min_lr_ratio", c.min_lr_ratio);
106	c.warmup_steps = j.value("warmup_steps", c.warmup_steps);	156	c.warmup_steps = j.value("warmup_steps", c.warmup_steps);
107	c.max_steps = j.value("max_steps", c.max_steps);	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	c.beta1 = j.value("beta1", c.beta1);	163	c.beta1 = j.value("beta1", c.beta1);
109	c.beta2 = j.value("beta2", c.beta2);	164	c.beta2 = j.value("beta2", c.beta2);
110	c.eps = j.value("eps", c.eps);	165	c.eps = j.value("eps", c.eps);
				`@@ -114,10 +169,21 @@ inline void from_json(const nlohmann::json& j, TrainConfig& c) {`
114	c.grad_accum_steps = j.value("grad_accum_steps", c.grad_accum_steps);	169	c.grad_accum_steps = j.value("grad_accum_steps", c.grad_accum_steps);
115	c.precision = j.value("precision", c.precision);	170	c.precision = j.value("precision", c.precision);
116	c.checkpoint_every = j.value("checkpoint_every", c.checkpoint_every);	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	c.eval_every = j.value("eval_every", c.eval_every);	174	c.eval_every = j.value("eval_every", c.eval_every);
118	c.eval_batches = j.value("eval_batches", c.eval_batches);	175	c.eval_batches = j.value("eval_batches", c.eval_batches);
119	c.seed = j.value("seed", c.seed);	176	c.seed = j.value("seed", c.seed);
120	c.deterministic = j.value("deterministic", c.deterministic);	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	inline Config load_config(const std::string& path) {	189	inline Config load_config(const std::string& path) {

modified src/train/optimizer.cpp +173 −24

				`@@ -1,6 +1,7 @@`
1	// Author: Simon-Pierre Boucher — contact@spboucher.ai	1	// Author: Simon-Pierre Boucher — contact@spboucher.ai
2	#include "train/optimizer.h"	2	#include "train/optimizer.h"
3		3
		4	+#include "ops/cpu/cpu_ops.h"
4	#include "ops/metal/metal_ops.h"	5	#include "ops/metal/metal_ops.h"
5	#include "ops/ops.h"	6	#include "ops/ops.h"
6		7
				`@@ -9,7 +10,43 @@`
9		10
10	namespace forge::train {	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	: opts_(opts) {	50	: opts_(opts) {
14	std::unordered_set<const void*> seen;	51	std::unordered_set<const void*> seen;
15	for (const auto& [name, p] : named_params) {	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	if (!seen.insert(p.id()).second) continue; // tied param, already tracked	54	if (!seen.insert(p.id()).second) continue; // tied param, already tracked
18	params_.push_back(p);	55	params_.push_back(p);
19	decay_.push_back(p.value().ndim() >= 2);	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	double sq = 0.0;	62	double sq = 0.0;
25	for (const Var& p : params_) {	63	for (const Var& p : params_) {
26	if (!p.has_grad()) continue;	64	if (!p.has_grad()) continue;
				`@@ -39,7 +77,7 @@ float AdamW::clip_global_norm(float max_norm) {`
39	return norm;	77	return norm;
40	}	78	}
41		79
42	−void AdamW::ensure_state() {	80	+void Optimizer::ensure_state() {
43	if (!m_.empty()) return;	81	if (!m_.empty()) return;
44	m_.reserve(params_.size());	82	m_.reserve(params_.size());
45	v_.reserve(params_.size());	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	if (ops::backend() == ops::Backend::CPU) {	91	if (ops::backend() == ops::Backend::CPU) {
54	const float norm = clip_global_norm(max_norm);	92	const float norm = clip_global_norm(max_norm);
55	step(lr);	93	step(lr);
				`@@ -80,44 +118,155 @@ float AdamW::step_with_clip(float lr, float max_norm) {`
80	ensure_state(); // CPU-side zeros are safe here: the stream is closed	118	ensure_state(); // CPU-side zeros are safe here: the stream is closed
81	++t_;	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	metal::ConcurrentRegion region;	122	metal::ConcurrentRegion region;
85	for (size_t i = 0; i < params_.size(); ++i) {	123	for (size_t i = 0; i < params_.size(); ++i) {
		124	+ if (muon_[i]) continue;
86	const float wd = decay_[i] ? opts_.weight_decay : 0.0f;	125	const float wd = decay_[i] ? opts_.weight_decay : 0.0f;
87	metal::adamw_step(params_[i].value(), params_[i].grad(), m_[i], v_[i], lr,	126	metal::adamw_step(params_[i].value(), params_[i].grad(), m_[i], v_[i], lr,
88	opts_.beta1, opts_.beta2, t_, opts_.eps, wd, grad_scale);	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	metal::sync(); // weights final before the next step's CPU-side zero_grad	194	metal::sync(); // weights final before the next step's CPU-side zero_grad
92	return norm;	195	return norm;
93	}	196	}
94		197
95	−void AdamW::step(float lr) {	198	+void Optimizer::adamw_cpu(size_t pi, float lr) {
96	− ensure_state();
97	− ++t_; // 1-based: bias correction divides by (1 - beta^t)
98	const float bc1 = 1.0f - std::pow(opts_.beta1, float(t_));	199	const float bc1 = 1.0f - std::pow(opts_.beta1, float(t_));
99	const float bc2 = 1.0f - std::pow(opts_.beta2, 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	+}
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	for (size_t pi = 0; pi < params_.size(); ++pi) {	262	for (size_t pi = 0; pi < params_.size(); ++pi) {
102	− Var& p = params_[pi];	263	+ if (!params_[pi].has_grad()) continue;
103	− if (!p.has_grad()) continue;	264	+ if (muon_[pi]) muon_cpu(pi, lr);
104	− const float wd = decay_[pi] ? opts_.weight_decay : 0.0f;	265	+ else adamw_cpu(pi, lr);
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	− }
117	}	266	}
118	}	267	}
119		268
120	−void AdamW::zero_grad() {	269	+void Optimizer::zero_grad() {
121	for (const Var& p : params_) p.zero_grad();	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	namespace forge::train {	10	namespace forge::train {
11		11
12	−// AdamW with decoupled weight decay (llm.c/PyTorch convention: eps outside	12	+// Two optimizer kinds behind one interface:
13	−// sqrt, wd only on dim>=2 params, decay folded into the same update).	13	+// - "adamw": AdamW with decoupled weight decay on every parameter
14	−// Deduplicates tied parameters by Var::id().	14	+// (llm.c/PyTorch convention: eps outside sqrt, wd only on dim>=2 params).
15	−class AdamW {	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	public:	23	public:
17	struct Options {	24	struct Options {
		25	+ std::string kind = "adamw"; // "adamw" \| "muon"
18	float beta1 = 0.9f;	26	float beta1 = 0.9f;
19	float beta2 = 0.95f;	27	float beta2 = 0.95f;
20	float eps = 1e-8f;	28	float eps = 1e-8f;
21	float weight_decay = 0.1f;	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	// Global-norm gradient clip: returns the pre-clip norm and scales all	36	// Global-norm gradient clip: returns the pre-clip norm and scales all
27	// grads by min(1, max_norm/norm). No-op when max_norm <= 0. (CPU path.)	37	// grads by min(1, max_norm/norm). No-op when max_norm <= 0. (CPU path.)
				`@@ -30,11 +40,12 @@ public:`
30	void step(float lr);	40	void step(float lr);
31	void zero_grad();	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	// grad norm. On Metal it expects backward() already encoded: it encodes	44	// grad norm. On Metal it expects backward() already encoded: it encodes
35	// per-tensor sumsq, syncs (this is the step's loss-readback boundary),	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	46	+ // folds the clip factor into the update, and syncs again at the end.
37	− // syncs again after the update.	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	float step_with_clip(float lr, float max_norm);	49	float step_with_clip(float lr, float max_norm);
39		50
40	// Allocate m/v now (checkpoint loading needs the buffers to exist).	51	// Allocate m/v now (checkpoint loading needs the buffers to exist).
				`@@ -48,11 +59,20 @@ public:`
48	const std::vector<Var>& params() const { return params_; }	59	const std::vector<Var>& params() const { return params_; }
49		60
50	private:	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	Options opts_;	67	Options opts_;
52	std::vector<Var> params_;	68	std::vector<Var> params_;
53	std::vector<bool> decay_; // dim >= 2	69	std::vector<bool> decay_; // dim >= 2
		70	+ std::vector<bool> muon_; // kind=="muon" && 2-D && not embedding/head
54	std::vector<Tensor> m_, v_; // f32, allocated lazily at first step	71	std::vector<Tensor> m_, v_; // f32, allocated lazily at first step
55	int64_t t_ = 0;	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	} // namespace forge::train	78	} // namespace forge::train

modified src/train/scheduler.h +15 −0

+     return min_lr + coeff * (max_lr - min_lr);
+ }
++// Warmup–Stable–Decay (MiniCPM): linear warmup, flat plateau at max_lr, then a
++// short cooldown over the final decay_steps. Any plateau checkpoint can be
++// resumed and the run extended without re-deciding total_steps up front. The
++// cooldown uses the 1-sqrt shape, which beats linear/exponential cooldowns in
++// the WSD ablations (arXiv 2404.06395 follow-ups).
++inline float lr_wsd(int64_t step, float max_lr, float min_lr, int64_t warmup_steps,
++                    int64_t total_steps, int64_t decay_steps) {
++    if (step < warmup_steps) return max_lr * float(step + 1) / float(warmup_steps + 1);
++    const int64_t decay_start = total_steps - decay_steps;
++    if (step < decay_start) return max_lr;
++    if (step >= total_steps) return min_lr;
++    const float ratio = float(step - decay_start) / float(decay_steps);
++    return min_lr + (max_lr - min_lr) * (1.0f - std::sqrt(ratio));
++}
++
+ } // namespace forge::train

modified src/train/trainer.cpp +33 −5

				`@@ -1,6 +1,7 @@`
1	// Author: Simon-Pierre Boucher — contact@spboucher.ai	1	// Author: Simon-Pierre Boucher — contact@spboucher.ai
2	#include "train/trainer.h"	2	#include "train/trainer.h"
3		3
		4	+#include "core/fmodel.h"
4	#include "ops/metal/metal_ops.h"	5	#include "ops/metal/metal_ops.h"
5	#include "ops/ops.h"	6	#include "ops/ops.h"
6	#include "train/checkpoint.h"	7	#include "train/checkpoint.h"
				`@@ -8,6 +9,7 @@`
8		9
9	#include <Foundation/Foundation.hpp>	10	#include <Foundation/Foundation.hpp>
10		11
		12	+#include <algorithm>
11	#include <chrono>	13	#include <chrono>
12	#include <cstdio>	14	#include <cstdio>
13	#include <filesystem>	15	#include <filesystem>
				`@@ -20,12 +22,15 @@ Trainer::Trainer(Config cfg, const std::string& data_dir, const std::string& out`
20	std::filesystem::create_directories(out_dir_);	22	std::filesystem::create_directories(out_dir_);
21	model_ = std::make_unique<nn::Transformer>(cfg_.model, cfg_.train.seed);	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	opts.beta1 = cfg_.train.beta1;	27	opts.beta1 = cfg_.train.beta1;
25	opts.beta2 = cfg_.train.beta2;	28	opts.beta2 = cfg_.train.beta2;
26	opts.eps = cfg_.train.eps;	29	opts.eps = cfg_.train.eps;
27	opts.weight_decay = cfg_.train.weight_decay;	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	train_data_ = std::make_unique<DataLoader>(data_dir + "/train.bin",	35	train_data_ = std::make_unique<DataLoader>(data_dir + "/train.bin",
31	cfg_.model.context_length,	36	cfg_.model.context_length,
				`@@ -64,6 +69,22 @@ void Trainer::save(int64_t step) {`
64	save_checkpoint(out_dir_ + "/ckpt_latest.bin", model_->named_parameters(), opt_.get(),	69	save_checkpoint(out_dir_ + "/ckpt_latest.bin", model_->named_parameters(), opt_.get(),
65	meta);	70	meta);
66	std::printf("checkpoint saved: %s\n", (out_dir_ + name).c_str());	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	void Trainer::train(const std::string& resume_from) {	90	void Trainer::train(const std::string& resume_from) {
				`@@ -87,18 +108,25 @@ void Trainer::train(const std::string& resume_from) {`
87	const int64_t min_lr_steps = tc.max_steps;	108	const int64_t min_lr_steps = tc.max_steps;
88	const float min_lr = tc.lr * tc.min_lr_ratio;	109	const float min_lr = tc.lr * tc.min_lr_ratio;
89	const int64_t tokens_per_step = B * T * tc.grad_accum_steps;	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	cfg_.model.name.c_str(),	117	cfg_.model.name.c_str(),
93	static_cast<long long>(cfg_.model.num_params()),	118	static_cast<long long>(cfg_.model.num_params()),
94	static_cast<long long>(tc.max_steps),	119	static_cast<long long>(tc.max_steps),
95	static_cast<long long>(tokens_per_step),	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	for (int64_t step = start_step; step < tc.max_steps; ++step) {	124	for (int64_t step = start_step; step < tc.max_steps; ++step) {
99	NS::AutoreleasePool* pool = NS::AutoreleasePool::alloc()->init();	125	NS::AutoreleasePool* pool = NS::AutoreleasePool::alloc()->init();
100	const auto t0 = std::chrono::steady_clock::now();	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	opt_->zero_grad();	131	opt_->zero_grad();
104	const bool on_gpu = ops::backend() == ops::Backend::Metal;	132	const bool on_gpu = ops::backend() == ops::Backend::Metal;
105		133

spb/forge

Add configurable training modes: Muon optimizer and WSD schedule

Theme