diff --git a/baseline/experiments/stock_gpt2_sgd_momentum/README.md b/baseline/experiments/stock_gpt2_sgd_momentum/README.md new file mode 100644 index 00000000..477e24cc --- /dev/null +++ b/baseline/experiments/stock_gpt2_sgd_momentum/README.md @@ -0,0 +1,29 @@ +# Stock GPT-2 Small / FineWeb10B — classical SGD + momentum + +This experiment reuses the **unchanged, SHA256-verified upstream GPT-2 Small model** and the same pinned FineWeb10B data and training protocol as the AdamW/MuonClip control. No rotary embeddings, RMSNorm, six-head attention, U-Net skips, or untied head. + +Architecture: 12 layers, 12 heads, width 768, MLP 3072, context 1024, vocabulary 50257, 124439808 unique parameters; packed QKV, learned positions, affine LayerNorm, GPT-2 GELU, tied token/output embeddings. + +## Candidate SGD recipe (not yet empirically optimized) + +- Optimizer: `torch.optim.SGD`, classical momentum **0.90**, Nesterov **off**, dampening 0, weight decay **0** on all parameters. +- Peak learning rate **0.01** is an initial screening candidate, **not a proven best value**. Do not label it optimal until a validation-only sweep on the actual GPT-2 model/data/hardware is complete. +- 700 warmup updates, cosine decay to zero; global gradient norm clipping at 1.0. +- 19560 optimizer updates, 524288 tokens/update, microbatch 64 sequences with 8 accumulation steps. +- Same FineWeb10B revision, fixed validation tokens, seed, spectral checkpoints, and stopping budget as the existing controls. +- Zero weight decay intentionally isolates SGD + momentum and avoids additional weight-scale effects; it differs from the AdamW control's matrix weight decay. + +## Run + +```bash +python3 -m unittest baseline/gpt2_small/muon_speedrun/test_sgd_momentum.py +python3 baseline/gpt2_small/muon_speedrun/cloudshell.py start --node YOUR_TPU_NODE --optimizer sgd_momentum --hours 12 +``` + +The launcher checks the existing TPU lease and performs the full-model TPU forward/backward/optimizer preflight before training. No TPU allocation or training is initiated by this change. + +## Qualification before calling these hyperparameters best + +On the fixed architecture/data/batch/seed/horizon, sweep peak LR `[0.001, 0.003, 0.01, 0.03, 0.1]` and momentum `[0.9, 0.95]`; reject nonfinite/divergent runs, rank by held-out **validation NLL only**, then repeat the leading candidates over three seeds and a full 19560-update horizon. Preserve the untouched test set until the recipe is locked. These sweeps are a **proposed protocol**, not completed experiments. + +Compare loss, perplexity and WeightWatcher alpha against the **stock-GPT-2** controls only; historical modified-architecture speedruns are not comparable. diff --git a/baseline/gpt2_small/muon_speedrun/benchmark_config.py b/baseline/gpt2_small/muon_speedrun/benchmark_config.py index a76b0f80..13562e4f 100644 --- a/baseline/gpt2_small/muon_speedrun/benchmark_config.py +++ b/baseline/gpt2_small/muon_speedrun/benchmark_config.py @@ -12,6 +12,9 @@ VAL_TOKENS = 10485760 WARMUP = 700 PEAK_ADAM_LR = 0.0006 +# Untuned starting point: NOT an empirically established best SGD recipe. +SGD_PEAK_LR = 0.01 +SGD_MOMENTUM = 0.9 GRAD_CLIP = 1.0 MEASUREMENT_INTERVAL = 250 @@ -30,7 +33,9 @@ def protocol(): source_commit=SOURCE_COMMIT, steps=TOTAL_STEPS, batch_tokens=BATCH_TOKENS, tokens=TOTAL_STEPS * BATCH_TOKENS, warmup_updates=WARMUP, lr_schedule='cosine_to_zero', - adam_peak_lr=PEAK_ADAM_LR, adam_betas=[0.9, 0.95], + adam_peak_lr=PEAK_ADAM_LR, + sgd_peak_lr=SGD_PEAK_LR, sgd_momentum=SGD_MOMENTUM, + sgd_weight_decay=0.0, sgd_nesterov=False, adam_betas=[0.9, 0.95], adam_eps=1e-8, matrix_weight_decay=0.1, gradient_clip_norm=GRAD_CLIP, measurement_interval=MEASUREMENT_INTERVAL, diff --git a/baseline/gpt2_small/muon_speedrun/cloudshell.py b/baseline/gpt2_small/muon_speedrun/cloudshell.py index 6394529a..99dfded8 100644 --- a/baseline/gpt2_small/muon_speedrun/cloudshell.py +++ b/baseline/gpt2_small/muon_speedrun/cloudshell.py @@ -184,7 +184,7 @@ def main(): p = argparse.ArgumentParser(description=__doc__) p.add_argument('action',choices=('start','status')) p.add_argument('--hours',type=float,default=12) - p.add_argument('--optimizer',choices=('muon','muon_clip','adam','adamw'),default='muon_clip') + p.add_argument('--optimizer',choices=('muon','muon_clip','adam','adamw','sgd_momentum'),default='muon_clip') p.add_argument('--microbatch',type=int,choices=(32,64,128),default=64) p.add_argument('--attention',choices=('auto','flash','math'),default='flash') p.add_argument('--replace-current',action='store_true',help='Stop the previous speedrun and start from initialization') diff --git a/baseline/gpt2_small/muon_speedrun/optim.py b/baseline/gpt2_small/muon_speedrun/optim.py index d02250b4..0015a60a 100644 --- a/baseline/gpt2_small/muon_speedrun/optim.py +++ b/baseline/gpt2_small/muon_speedrun/optim.py @@ -3,7 +3,7 @@ import math import torch from stock_config import ARCHITECTURE -from benchmark_config import BENCHMARK, lr_factor, GRAD_CLIP +from benchmark_config import BENCHMARK, lr_factor, GRAD_CLIP, SGD_PEAK_LR, SGD_MOMENTUM def schedule(step): @@ -80,7 +80,7 @@ def load_state_dict(self, state): def make_optimizers(model, rt, kind='muon'): - if kind not in ('muon', 'muon_clip', 'adam', 'adamw'): + if kind not in ('muon', 'muon_clip', 'adam', 'adamw', 'sgd_momentum'): raise ValueError('Unknown optimizer: '+kind) if getattr(model, 'architecture_id', None) == ARCHITECTURE: return make_stock_optimizers(model, rt, kind) @@ -119,6 +119,17 @@ def make_stock_optimizers(model, rt, kind): """ matrices = [(n, p) for n, p in model.transformer.h.named_parameters() if p.ndim == 2] muon = Muon(matrices, rt) if kind == 'muon' else None + if kind == 'sgd_momentum': + # Plain classical momentum: no Adam, Nesterov, adaptive rescaling or decay. + # Every unique trainable parameter, including tied embeddings, has one owner. + groups = [dict(params=[p for p in model.parameters() if p.ndim >= 2], role='matrices', + lr=SGD_PEAK_LR, peak_lr=SGD_PEAK_LR, weight_decay=0.0), + dict(params=[p for p in model.parameters() if p.ndim < 2], role='bias_and_layernorm', + lr=SGD_PEAK_LR, peak_lr=SGD_PEAK_LR, weight_decay=0.0)] + sgd = torch.optim.SGD(groups, momentum=SGD_MOMENTUM, nesterov=False, + dampening=0.0, foreach=False) + sgd.benchmark = BENCHMARK + return None, sgd if kind == 'muon_clip': from muonclip import MuonClip muon = MuonClip(matrices, rt, model) @@ -138,6 +149,11 @@ def make_stock_optimizers(model, rt, kind): def optimizer_metadata(adam): + if isinstance(adam, torch.optim.SGD): + return {'class':'torch.optim.SGD', 'groups':[ + {'role':g['role'], 'peak_lr':g['peak_lr'], 'weight_decay':g['weight_decay'], + 'momentum':g['momentum'], 'nesterov':g['nesterov'], + 'parameters':sum(p.numel() for p in g['params'])} for g in adam.param_groups]} return {'class':'torch.optim.'+type(adam).__name__, 'groups':[{'role':g['role'], 'peak_lr':g['peak_lr'], 'weight_decay':g['weight_decay'], 'betas':list(g['betas']), diff --git a/baseline/gpt2_small/muon_speedrun/run.py b/baseline/gpt2_small/muon_speedrun/run.py index 4063b06a..82424df8 100644 --- a/baseline/gpt2_small/muon_speedrun/run.py +++ b/baseline/gpt2_small/muon_speedrun/run.py @@ -214,12 +214,13 @@ def train(a): 'head_tied_to_embedding':True, 'dropout':0.0, 'adam_control_matrix_lr':0.0006 if not muon else None, 'warmup_updates':WARMUP, 'lr_schedule':'cosine_to_zero', - 'weight_decay':0.0 if a.optimizer == 'adam' else 0.1, + 'weight_decay':0.0 if a.optimizer in ('adam', 'sgd_momentum') else 0.1, 'weight_decay_scope':('auxiliary_matrices_only' if a.optimizer == 'muon' else 'all_matrices') - if a.optimizer != 'adam' else 'none', + if a.optimizer not in ('adam', 'sgd_momentum') else 'none', 'muon_weight_decay':getattr(muon, 'weight_decay', 0.0) if muon else None, 'muonclip':dict(threshold=muon.threshold, balance=muon.balance, rms_scale=muon.rms_scale, momentum=muon.beta) if a.optimizer == 'muon_clip' else None, - 'adam_implementation':optimizer_metadata(adam), + 'optimizer_implementation':optimizer_metadata(adam), + 'adam_implementation':optimizer_metadata(adam) if a.optimizer != 'sgd_momentum' else None, 'historical_reference':'muon-speedrun-muon-20261005-030026', 'gradient_clipping':True, 'gradient_clip_norm':GRAD_CLIP, 'attention':a.attention, 'data_repo':source.manifest['repo'], 'data_revision':source.manifest['revision'], @@ -321,7 +322,7 @@ def main(): p.add_argument('--device', choices=('tpu', 'cpu'), default='tpu') p.add_argument('--attention', choices=('flash', 'math'), default='flash') p.add_argument('--microbatch', type=int, choices=(32,64,128), default=64) - p.add_argument('--optimizer', choices=('muon', 'muon_clip', 'adam', 'adamw'), default='muon_clip') + p.add_argument('--optimizer', choices=('muon', 'muon_clip', 'adam', 'adamw', 'sgd_momentum'), default='muon_clip') p.add_argument('--seed', type=int, default=42) budget = p.add_mutually_exclusive_group() budget.add_argument('--full-budget', action='store_true', default=True, help='Default: complete all 19560 updates') diff --git a/baseline/gpt2_small/muon_speedrun/test_sgd_momentum.py b/baseline/gpt2_small/muon_speedrun/test_sgd_momentum.py new file mode 100644 index 00000000..d07a71c5 --- /dev/null +++ b/baseline/gpt2_small/muon_speedrun/test_sgd_momentum.py @@ -0,0 +1,52 @@ +"""CPU regression checks for the stock-GPT-2 momentum-SGD optimizer arm.""" +import sys +import unittest +from pathlib import Path +from types import SimpleNamespace + +import torch + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +from optim import make_stock_optimizers, apply_update, optimizer_metadata +from benchmark_config import SGD_PEAK_LR, SGD_MOMENTUM, lr_factor + + +class TinyTiedModel(torch.nn.Module): + architecture_id = "not-needed-for-direct-stock-factory" + + def __init__(self): + super().__init__() + self.transformer = torch.nn.Module() + self.transformer.h = torch.nn.ModuleList([torch.nn.Linear(4, 4)]) + self.transformer.wte = torch.nn.Embedding(8, 4) + self.lm_head = torch.nn.Linear(4, 8, bias=False) + self.lm_head.weight = self.transformer.wte.weight + self.norm = torch.nn.LayerNorm(4) + + +class SGDTests(unittest.TestCase): + def test_tied_parameter_once_and_classical_momentum(self): + model = TinyTiedModel() + muon, opt = make_stock_optimizers(model, SimpleNamespace(tpu=False), "sgd_momentum") + self.assertIsNone(muon) + self.assertIsInstance(opt, torch.optim.SGD) + params = [p for g in opt.param_groups for p in g["params"]] + self.assertEqual(len(params), len(set(map(id, params)))) + self.assertEqual(set(map(id, params)), set(map(id, model.parameters()))) + self.assertTrue(all(g["weight_decay"] == 0 for g in opt.param_groups)) + self.assertTrue(all(g["momentum"] == SGD_MOMENTUM for g in opt.param_groups)) + self.assertTrue(all(not g["nesterov"] for g in opt.param_groups)) + self.assertEqual(optimizer_metadata(opt)["class"], "torch.optim.SGD") + p = model.transformer.h[0].weight + old = p.detach().clone() + for param in model.parameters(): + param.grad = torch.ones_like(param) + apply_update(muon, opt, SimpleNamespace(scalar=lambda x: x), 0) + self.assertTrue(torch.allclose(p, old - SGD_PEAK_LR * lr_factor(0))) + previous = p.detach().clone() + apply_update(muon, opt, SimpleNamespace(scalar=lambda x: x), 1) + self.assertTrue(torch.allclose(p, previous - SGD_PEAK_LR * lr_factor(1) * (1 + SGD_MOMENTUM))) + + +if __name__ == "__main__": + unittest.main() diff --git a/baseline/gpt2_small/muon_speedrun/worker.py b/baseline/gpt2_small/muon_speedrun/worker.py index 7a08752e..38d70545 100644 --- a/baseline/gpt2_small/muon_speedrun/worker.py +++ b/baseline/gpt2_small/muon_speedrun/worker.py @@ -102,7 +102,7 @@ def main(): p = argparse.ArgumentParser() p.add_argument('root', type=Path) p.add_argument('deadline', type=float) - p.add_argument('--optimizer', choices=('muon','muon_clip','adam','adamw'), default='muon_clip') + p.add_argument('--optimizer', choices=('muon','muon_clip','adam','adamw','sgd_momentum'), default='muon_clip') p.add_argument('--seed', type=int, default=42) p.add_argument('--full-budget', action='store_true', default=True) p.add_argument('--stop-at-target', action='store_false', dest='full_budget')