diff --git a/baseline/experiments/stock_gpt2_sgd_momentum/README.md b/baseline/experiments/stock_gpt2_sgd_momentum/README.md index 289d4a6d..27226c34 100644 --- a/baseline/experiments/stock_gpt2_sgd_momentum/README.md +++ b/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. diff --git a/baseline/gpt2_small/muon_speedrun/benchmark_config.py b/baseline/gpt2_small/muon_speedrun/benchmark_config.py index b1a6394d..8e78c8a7 100644 --- a/baseline/gpt2_small/muon_speedrun/benchmark_config.py +++ b/baseline/gpt2_small/muon_speedrun/benchmark_config.py @@ -9,14 +9,14 @@ 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 @@ -24,7 +24,7 @@ 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): @@ -32,11 +32,6 @@ 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)) diff --git a/baseline/gpt2_small/muon_speedrun/cloudshell.py b/baseline/gpt2_small/muon_speedrun/cloudshell.py index 1f3a1fb7..a068d744 100644 --- a/baseline/gpt2_small/muon_speedrun/cloudshell.py +++ b/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():