From c2327cdcb25cdf8d8548ac97e31f1581684761ca Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Fri, 16 Oct 2026 17:07:44 -0700 Subject: [PATCH 1/6] Configure published-center SGDM short speedrun --- .../muon_speedrun/benchmark_config.py | 31 +++++++++++++------ 1 file changed, 22 insertions(+), 9 deletions(-) diff --git a/baseline/gpt2_small/muon_speedrun/benchmark_config.py b/baseline/gpt2_small/muon_speedrun/benchmark_config.py index 13562e4f..b1a6394d 100644 --- a/baseline/gpt2_small/muon_speedrun/benchmark_config.py +++ b/baseline/gpt2_small/muon_speedrun/benchmark_config.py @@ -4,19 +4,27 @@ 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): @@ -24,18 +32,23 @@ def lr_factor(step): 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, From 254af590e4274715ffa33a8ba4d62eda9b3240c3 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Fri, 16 Oct 2026 17:07:46 -0700 Subject: [PATCH 2/6] Configure published-center SGDM short speedrun --- baseline/gpt2_small/muon_speedrun/optim.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/baseline/gpt2_small/muon_speedrun/optim.py b/baseline/gpt2_small/muon_speedrun/optim.py index 0015a60a..7749ca71 100644 --- a/baseline/gpt2_small/muon_speedrun/optim.py +++ b/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'], From cfcaa5e1772a216a275024056c45aef8df09b5e3 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Fri, 16 Oct 2026 17:07:48 -0700 Subject: [PATCH 3/6] Configure published-center SGDM short speedrun --- baseline/gpt2_small/muon_speedrun/run.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/baseline/gpt2_small/muon_speedrun/run.py b/baseline/gpt2_small/muon_speedrun/run.py index 82424df8..8405f3d2 100644 --- a/baseline/gpt2_small/muon_speedrun/run.py +++ b/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() From 8dd2ef6a74a5c1be1108338c4d0e2b13d7aa873c Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Fri, 16 Oct 2026 17:07:50 -0700 Subject: [PATCH 4/6] Configure published-center SGDM short speedrun --- baseline/gpt2_small/muon_speedrun/cloudshell.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/baseline/gpt2_small/muon_speedrun/cloudshell.py b/baseline/gpt2_small/muon_speedrun/cloudshell.py index 99dfded8..1f3a1fb7 100644 --- a/baseline/gpt2_small/muon_speedrun/cloudshell.py +++ b/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) From 5eea8cb39ae83292c9a31042bbbe01fc6725abbe Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Fri, 16 Oct 2026 17:07:52 -0700 Subject: [PATCH 5/6] Configure published-center SGDM short speedrun --- .../muon_speedrun/test_sgd_momentum.py | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/baseline/gpt2_small/muon_speedrun/test_sgd_momentum.py b/baseline/gpt2_small/muon_speedrun/test_sgd_momentum.py index d07a71c5..8c0ff3b9 100644 --- a/baseline/gpt2_small/muon_speedrun/test_sgd_momentum.py +++ b/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() From c4087d76b8434b6b47f07a6b5467e353f5ca71f1 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Fri, 16 Oct 2026 17:07:54 -0700 Subject: [PATCH 6/6] Configure published-center SGDM short speedrun --- baseline/experiments/stock_gpt2_sgd_momentum/README.md | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/baseline/experiments/stock_gpt2_sgd_momentum/README.md b/baseline/experiments/stock_gpt2_sgd_momentum/README.md index 477e24cc..289d4a6d 100644 --- a/baseline/experiments/stock_gpt2_sgd_momentum/README.md +++ b/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.