Add classical SGD + momentum to stock GPT-2 Small FineWeb TPU speedrun

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.

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,

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')

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']),

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')

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()

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')