Run stock GPT-2 SGDM to validation target on full horizon

baseline/experiments/stock_gpt2_sgd_momentum/README.md

@@ -1,29 +1,31 @@
-# Stock GPT-2 Small / FineWeb10B — classical SGD + momentum
+# Stock GPT-2 Small / FineWeb10B — SGD + momentum target run
 
-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.
+This experiment uses the unchanged SHA256-verified upstream GPT-2 Small model: 12 layers, 12 heads, width 768, MLP 3072, context 1024, vocabulary 50257, and 124439808 unique parameters.
 
-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.
+## Optimizer
 
-## Published-center short speedrun profile
+- Classical `torch.optim.SGD`
+- peak LR **0.3**
+- momentum **0.90**
+- Nesterov off
+- matrix/embedding weight decay **3e-5**
+- bias/LayerNorm weight decay 0
+- global gradient clipping 1.0
 
-- 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.
+## Full target-reaching schedule
 
-## Run
+The previous 6200-update short profile cooled down before the model was converged. It is retained only as historical output and must not be treated as a converged SGD control.
 
-```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 replacement run starts from fresh initialization and uses the full original GPT-2/FineWeb horizon: **19560 updates**, **700-update linear warmup**, then **cosine decay to zero**. Validation and WeightWatcher run every **250 updates** and at exit. With `--stop-at-target`, training stops only after a full 10485760-token validation reaches **NLL <= 3.28**; otherwise the 19560-update schedule is the safety cap.
 
-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.
+This combines the published SGDM optimizer center (LR 0.3, momentum 0.9, small matrix decay) with the repository's original full GPT-2/FineWeb training horizon. It is not claimed to be a published end-to-end SGDM schedule.
 
-## Qualification before calling these hyperparameters best
+## Launch
 
-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.
+```bash
+python3 baseline/gpt2_small/muon_speedrun/cloudshell.py start \
+  --node YOUR_TPU_NODE --optimizer sgd_momentum --hours 12 \
+  --stop-at-target --replace-current
+```
 
-Compare loss, perplexity and WeightWatcher alpha against the **stock-GPT-2** controls only; historical modified-architecture speedruns are not comparable.
+The launcher performs dependency, attention, and full-model optimizer preflights before training. The target is selected solely by validation NLL; alpha is diagnostic and is never used for stopping.

baseline/gpt2_small/muon_speedrun/benchmark_config.py

@@ -9,34 +9,29 @@
 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')
+BENCHMARK = ('stock-gpt2-124m-fineweb10b-sgdm-full-target-v1' if PROFILE == SGD_SPEEDRUN_PROFILE else 'llmc-gpt2-124m-fineweb10b-2024-10-13-upstream-tpu-v2')
 SOURCE_COMMIT = '7ecd8906afe6ed7a2b2cdb731c042f26d525b820'
-TOTAL_STEPS = 6200 if PROFILE == SGD_SPEEDRUN_PROFILE else 19560
+TOTAL_STEPS = 19560
 BATCH_TOKENS = 524288
 VAL_TOKENS = 10485760
-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'
+WARMUP = 700
+WARMDOWN = 0
+LR_SCHEDULE = 'cosine_to_zero'
 PEAK_ADAM_LR = 0.0006
 # 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 = 125 if PROFILE == SGD_SPEEDRUN_PROFILE else 250
+MEASUREMENT_INTERVAL = 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))
 

baseline/gpt2_small/muon_speedrun/cloudshell.py

@@ -96,7 +96,7 @@ def stop_longrun():
 
 
 def start_remote(commit, hours=12, optimizer='muon_clip', microbatch=64, attention='flash', replace_current=False,
-                 launch_id=None, replace_longrun=False, checked_lease=None):
+                 launch_id=None, replace_longrun=False, checked_lease=None, stop_at_target=False):
     if os.geteuid() != 0 or not os.path.ismount('/mnt/disks/rg-data'):
         raise RuntimeError('Requires the existing mounted disk and root')
     if not re.fullmatch('[0-9a-f]{40}',commit):
@@ -160,7 +160,7 @@ def start_remote(commit, hours=12, optimizer='muon_clip', microbatch=64, attenti
                   'started_unix':time.time(),'deadline_unix':deadline,'hours_cap':hours,
                   'target_val_nll':3.28,'checkpoint_interval':MEASUREMENT_INTERVAL,'microbatch':microbatch,
                   'weightwatcher_interval':MEASUREMENT_INTERVAL, 'validation_token_error':True,
-                  'fresh_initialization':True, 'automatic_restart':False,
+                  'fresh_initialization':True, 'automatic_restart':False, 'stop_at_target':stop_at_target,
                   'historical_reference':'muon-speedrun-muon-20261005-030026',
                   'experiment':EXPERIMENT, 'cloud_uri':cloud_uri(root)}
         (root/'launch.json').write_text(json.dumps(record,indent=2))
@@ -174,14 +174,16 @@ def start_remote(commit, hours=12, optimizer='muon_clip', microbatch=64, attenti
         command += ['/mnt/disks/rg-data/continuous8/venv/bin/python','-u',
                     str(base/'muon_speedrun/worker.py'),str(root),str(deadline),
                     '--optimizer',optimizer,'--microbatch',str(microbatch),'--attention',attention]
+        if stop_at_target:
+            command += ['--stop-at-target']
         run(command)
         temp = LATEST.with_suffix('.tmp')
         temp.write_text(json.dumps(record,indent=2))
         temp.replace(LATEST)
         print('Started '+optimizer+' recipe: '+unit,flush=True)
         print('Log: '+str(root/'run.log'),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(f'{TOTAL_STEPS:,} update cap; profile {PROFILE}; stop_at_target={stop_at_target}.',flush=True)
+        print(f'Paired validation token error and raw/clipped WeightWatcher alpha every {MEASUREMENT_INTERVAL} 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)
 
@@ -195,6 +197,7 @@ def main():
     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')
     p.add_argument('--replace-longrun',action='store_true',help='Stop the recorded 25k training service, retaining its files')
+    p.add_argument('--stop-at-target',action='store_true',help='Stop after a full validation reaches NLL <= 3.28')
     p.add_argument('--node', default=NODE, help='Existing TPU node; the old October 4 allocation has expired')
     p.add_argument('--zone', default=ZONE)
     p.add_argument('--lease', help=argparse.SUPPRESS)
@@ -207,7 +210,7 @@ def main():
         raise ValueError('Hours must be between 1 and 12, bounded by existing allocation')
     if a.on_tpu:
         if a.action == 'start':
-            start_remote(a.commit,a.hours,a.optimizer,a.microbatch,a.attention,a.replace_current,a.launch_id,a.replace_longrun,a.lease)
+            start_remote(a.commit,a.hours,a.optimizer,a.microbatch,a.attention,a.replace_current,a.launch_id,a.replace_longrun,a.lease,a.stop_at_target)
         else:
             status_remote()
         return 0
@@ -218,6 +221,8 @@ def main():
         command += ['--replace-current']
     if a.replace_longrun:
         command += ['--replace-longrun']
+    if a.stop_at_target:
+        command += ['--stop-at-target']
     if a.action == 'start':
         repo = Path(__file__).resolve().parents[3]
         if run(['git','-C',str(repo),'status','--porcelain'],capture_output=True).stdout.strip():