Configure stock GPT-2 SGDM short speedrun

baseline/experiments/stock_gpt2_sgd_momentum/README.md

@@ -4,12 +4,12 @@ This experiment reuses the **unchanged, SHA256-verified upstream GPT-2 Small mod
 
 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)
+## Published-center short speedrun profile
 
-- 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.
+- Optimizer: `torch.optim.SGD`, classical momentum **0.90**, Nesterov **off**, dampening 0, matrix/embedding weight decay **3e-5**.
+- Peak learning rate **0.3**, matching the released GPT-2 Small SGDM experiment center.
+- 400 linear warmup updates, constant peak LR through update 4400, then 1800-update linear warmdown to zero; global gradient norm clipping at 1.0.
+- 6200 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.
 

baseline/gpt2_small/muon_speedrun/benchmark_config.py

@@ -4,38 +4,51 @@
 TPU port differences; this identifier must accompany every new comparison.
 """
 import math
+import os
 
-BENCHMARK = 'llmc-gpt2-124m-fineweb10b-2024-10-13-upstream-tpu-v2'
+LLMC_PROFILE = 'llmc_reference'
+SGD_SPEEDRUN_PROFILE = 'sgd_speedrun'
+PROFILE = os.environ.get('RG_GPT2_PROFILE', LLMC_PROFILE)
+BENCHMARK = ('stock-gpt2-124m-fineweb10b-sgdm-published-center-short-v1' if PROFILE == SGD_SPEEDRUN_PROFILE else 'llmc-gpt2-124m-fineweb10b-2024-10-13-upstream-tpu-v2')
 SOURCE_COMMIT = '7ecd8906afe6ed7a2b2cdb731c042f26d525b820'
-TOTAL_STEPS = 19560
+TOTAL_STEPS = 6200 if PROFILE == SGD_SPEEDRUN_PROFILE else 19560
 BATCH_TOKENS = 524288
 VAL_TOKENS = 10485760
-WARMUP = 700
+WARMUP = 400 if PROFILE == SGD_SPEEDRUN_PROFILE else 700
+WARMDOWN = 1800 if PROFILE == SGD_SPEEDRUN_PROFILE else 0
+LR_SCHEDULE = 'linear_warmup_constant_linear_warmdown' if PROFILE == SGD_SPEEDRUN_PROFILE else 'cosine_to_zero'
 PEAK_ADAM_LR = 0.0006
-# Untuned starting point: NOT an empirically established best SGD recipe.
-SGD_PEAK_LR = 0.01
+# Published GPT-2 Small SGDM center from tml-epfl/why-weight-decay.
+SGD_PEAK_LR = 0.3
 SGD_MOMENTUM = 0.9
+SGD_WEIGHT_DECAY = 3e-5
+SGD_NESTEROV = False
 GRAD_CLIP = 1.0
-MEASUREMENT_INTERVAL = 250
+MEASUREMENT_INTERVAL = 125 if PROFILE == SGD_SPEEDRUN_PROFILE else 250
 
 
 def lr_factor(step):
     if step < 0:
         raise ValueError('Training update must be nonnegative')
     if step < WARMUP:
         return (step + 1) / WARMUP
+    if PROFILE == SGD_SPEEDRUN_PROFILE:
+        warmdown_start = TOTAL_STEPS - WARMDOWN
+        if step < warmdown_start:
+            return 1.0
+        return max(0.0, (TOTAL_STEPS - step) / WARMDOWN)
     ratio = min(1.0, (step - WARMUP) / (TOTAL_STEPS - WARMUP))
     return 0.5 * (1.0 + math.cos(math.pi * ratio))
 
 
 def protocol():
-    return dict(benchmark=BENCHMARK, source_repo='karpathy/llm.c',
+    return dict(profile=PROFILE, benchmark=BENCHMARK, source_repo='karpathy/llm.c',
                 source_commit=SOURCE_COMMIT, steps=TOTAL_STEPS,
                 batch_tokens=BATCH_TOKENS, tokens=TOTAL_STEPS * BATCH_TOKENS,
-                warmup_updates=WARMUP, lr_schedule='cosine_to_zero',
+                warmup_updates=WARMUP, warmdown_updates=WARMDOWN, lr_schedule=LR_SCHEDULE,
                 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],
+                sgd_weight_decay=SGD_WEIGHT_DECAY, sgd_nesterov=SGD_NESTEROV, 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

@@ -19,7 +19,8 @@
 NODE = QUEUE+'-node'
 BASE = Path('/mnt/disks/rg-data/gpt2small')
 LATEST = BASE/'MUON_SPEEDRUN_LATEST.json'
-EXPERIMENT = 'stock_gpt2_fineweb_muonclip_adamw'
+PROFILE = os.environ.get('RG_GPT2_PROFILE', 'llmc_reference')
+EXPERIMENT = ('stock_gpt2_sgd_momentum' if PROFILE == 'sgd_speedrun' else 'stock_gpt2_fineweb_muonclip_adamw')
 CLOUD_BASE = 'gs://tpu-builders-504820-ww-continuous8/gpt2small'
 
 
@@ -100,6 +101,11 @@ def start_remote(commit, hours=12, optimizer='muon_clip', microbatch=64, attenti
         raise RuntimeError('Requires the existing mounted disk and root')
     if not re.fullmatch('[0-9a-f]{40}',commit):
         raise ValueError('Expected pinned commit')
+    global PROFILE, EXPERIMENT
+    profile = 'sgd_speedrun' if optimizer == 'sgd_momentum' else 'llmc_reference'
+    os.environ['RG_GPT2_PROFILE'] = profile
+    PROFILE = profile
+    EXPERIMENT = ('stock_gpt2_sgd_momentum' if profile == 'sgd_speedrun' else 'stock_gpt2_fineweb_muonclip_adamw')
     with (BASE/'port-check-launch.lock').open('a') as lock:
         fcntl.flock(lock,fcntl.LOCK_EX|fcntl.LOCK_NB)
         if LATEST.exists() and launch_id and json.loads(LATEST.read_text()).get('launch_id') == launch_id:
@@ -147,7 +153,7 @@ def start_remote(commit, hours=12, optimizer='muon_clip', microbatch=64, attenti
         env = {'PYTHONPATH':str(base/'src')+':'+str(base.parent/'nanogpt_one_head/src'),
                'PJRT_DEVICE':'TPU','TPU_ACCELERATOR_TYPE':'v5litepod-8',
                'OMP_NUM_THREADS':'4','OPENBLAS_NUM_THREADS':'4','MKL_NUM_THREADS':'4',
-               'TOKENIZERS_PARALLELISM':'false'}
+               'TOKENIZERS_PARALLELISM':'false', 'RG_GPT2_PROFILE':profile}
         record = {'root':str(root),'unit':unit,'commit':commit,'optimizer':optimizer,
                   'architecture':ARCHITECTURE, 'benchmark':BENCHMARK, 'steps':TOTAL_STEPS,
                   'launch_id':launch_id, 'lease':checked,
@@ -174,7 +180,7 @@ def start_remote(commit, hours=12, optimizer='muon_clip', microbatch=64, attenti
         temp.replace(LATEST)
         print('Started '+optimizer+' recipe: '+unit,flush=True)
         print('Log: '+str(root/'run.log'),flush=True)
-        print('19,560 updates, 700 warmup, cosine decay, gradient clipping 1.0; full-budget default.',flush=True)
+        print(f'{TOTAL_STEPS:,} updates; profile {PROFILE}; full-budget default.',flush=True)
         print('Paired validation token error and raw/clipped WeightWatcher alpha every 250 updates and final.',flush=True)
         print('Hard cutoff UTC: '+dt.datetime.fromtimestamp(deadline,dt.timezone.utc).isoformat(),flush=True)
         print('No automatic restart or new TPU allocation. The selected hours limit is a cap, not an ETA. A deadline-stopped run is incomplete.',flush=True)

baseline/gpt2_small/muon_speedrun/optim.py

@@ -3,7 +3,8 @@
 import math
 import torch
 from stock_config import ARCHITECTURE
-from benchmark_config import BENCHMARK, lr_factor, GRAD_CLIP, SGD_PEAK_LR, SGD_MOMENTUM
+from benchmark_config import (BENCHMARK, PROFILE, SGD_SPEEDRUN_PROFILE, lr_factor, GRAD_CLIP,
+                              SGD_PEAK_LR, SGD_MOMENTUM, SGD_WEIGHT_DECAY, SGD_NESTEROV)
 
 
 def schedule(step):
@@ -120,13 +121,13 @@ 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.
+        if PROFILE != SGD_SPEEDRUN_PROFILE:
+            raise RuntimeError('sgd_momentum requires the launcher-selected sgd_speedrun profile')
         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),
+                       lr=SGD_PEAK_LR, peak_lr=SGD_PEAK_LR, weight_decay=SGD_WEIGHT_DECAY),
                   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,
+        sgd = torch.optim.SGD(groups, momentum=SGD_MOMENTUM, nesterov=SGD_NESTEROV,
                               dampening=0.0, foreach=False)
         sgd.benchmark = BENCHMARK
         return None, sgd
@@ -152,7 +153,7 @@ 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'],
+             'momentum':g['momentum'], 'nesterov':g['nesterov'], 'dampening':g['dampening'],
              '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'],

baseline/gpt2_small/muon_speedrun/run.py

@@ -19,7 +19,8 @@
 from tracking import queue_snapshot
 
 from benchmark_config import (BENCHMARK, TOTAL_STEPS, BATCH_TOKENS, VAL_TOKENS,
-                              MEASUREMENT_INTERVAL, WARMUP, GRAD_CLIP,
+                              MEASUREMENT_INTERVAL, WARMUP, WARMDOWN, LR_SCHEDULE, GRAD_CLIP,
+                              SGD_PEAK_LR, SGD_MOMENTUM, SGD_WEIGHT_DECAY, SGD_NESTEROV,
                               lr_factor, protocol)
 
 TARGET = 3.28
@@ -135,7 +136,7 @@ def preflight_identity(a):
                 config=asdict(architecture.GPTConfig()),
                 source_sha256=architecture.SOURCE_SHA256,
                 optimizer=a.optimizer, microbatch=a.microbatch, device=a.device,
-                attention=a.attention, seed=a.seed, benchmark=BENCHMARK)
+                attention=a.attention, seed=a.seed, benchmark=BENCHMARK, protocol=protocol())
 
 
 def require_preflight(a):
@@ -213,12 +214,14 @@ def train(a):
                 'adam_head_lr':0.0006, 'adam_scalar_lr':0.0006,
                 '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 in ('adam', 'sgd_momentum') else 0.1,
-                'weight_decay_scope':('auxiliary_matrices_only' if a.optimizer == 'muon' else 'all_matrices')
-                                     if a.optimizer not in ('adam', 'sgd_momentum') else 'none',
+                'warmup_updates':WARMUP, 'warmdown_updates':WARMDOWN, 'lr_schedule':LR_SCHEDULE,
+                'weight_decay':SGD_WEIGHT_DECAY if a.optimizer == 'sgd_momentum' else (0.0 if a.optimizer == 'adam' else 0.1),
+                'weight_decay_scope':'matrices_only' if a.optimizer == 'sgd_momentum' else (('auxiliary_matrices_only' if a.optimizer == 'muon' else 'all_matrices') if a.optimizer != 'adam' 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,
+                'sgd_peak_lr':SGD_PEAK_LR if a.optimizer == 'sgd_momentum' else None,
+                'sgd_momentum':SGD_MOMENTUM if a.optimizer == 'sgd_momentum' else None,
+                'sgd_nesterov':SGD_NESTEROV if a.optimizer == 'sgd_momentum' else None,
                 'optimizer_implementation':optimizer_metadata(adam),
                 'adam_implementation':optimizer_metadata(adam) if a.optimizer != 'sgd_momentum' else None,
                 'historical_reference':'muon-speedrun-muon-20261005-030026',
@@ -325,7 +328,7 @@ def main():
     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')
+    budget.add_argument('--full-budget', action='store_true', default=True, help=f'Default: complete all {TOTAL_STEPS} updates')
     budget.add_argument('--stop-at-target', action='store_false', dest='full_budget', help='Explicit non-default early stop at NLL <= 3.28')
     p.add_argument('--legacy-attention-check', action='store_true', help=argparse.SUPPRESS)
     a = p.parse_args()

baseline/gpt2_small/muon_speedrun/test_sgd_momentum.py

@@ -1,14 +1,17 @@
 """CPU regression checks for the stock-GPT-2 momentum-SGD optimizer arm."""
+import os
 import sys
 import unittest
 from pathlib import Path
 from types import SimpleNamespace
 
+os.environ['RG_GPT2_PROFILE'] = 'sgd_speedrun'
+
 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
+from benchmark_config import SGD_PEAK_LR, SGD_MOMENTUM, SGD_WEIGHT_DECAY, TOTAL_STEPS, WARMUP, WARMDOWN, lr_factor
 
 
 class TinyTiedModel(torch.nn.Module):
@@ -33,11 +36,13 @@ def test_tied_parameter_once_and_classical_momentum(self):
         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))
+        groups = {g['role']:g for g in opt.param_groups}
+        self.assertEqual(groups['matrices']['weight_decay'], SGD_WEIGHT_DECAY)
+        self.assertEqual(groups['bias_and_layernorm']['weight_decay'], 0.0)
         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
+        p = model.norm.bias
         old = p.detach().clone()
         for param in model.parameters():
             param.grad = torch.ones_like(param)
@@ -48,5 +53,13 @@ def test_tied_parameter_once_and_classical_momentum(self):
         self.assertTrue(torch.allclose(p, previous - SGD_PEAK_LR * lr_factor(1) * (1 + SGD_MOMENTUM)))
 
 
+    def test_speedrun_schedule(self):
+        self.assertEqual(TOTAL_STEPS, 6200)
+        self.assertAlmostEqual(lr_factor(0), 1/WARMUP)
+        self.assertAlmostEqual(lr_factor(WARMUP-1), 1.0)
+        self.assertAlmostEqual(lr_factor(TOTAL_STEPS-WARMDOWN), 1.0)
+        self.assertEqual(lr_factor(TOTAL_STEPS), 0.0)
+
+
 if __name__ == "__main__":
     unittest.main()