From 462494699f215cbaf2ba155c7c2d5506c987e3a9 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Sat, 17 Oct 2026 17:01:01 -0700 Subject: [PATCH 1/5] Configure layerwise SGDM schedule --- .../gpt2_small/muon_speedrun/benchmark_config.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/baseline/gpt2_small/muon_speedrun/benchmark_config.py b/baseline/gpt2_small/muon_speedrun/benchmark_config.py index 8e78c8a..cb8ed3c 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-full-target-v1' if PROFILE == SGD_SPEEDRUN_PROFILE else 'llmc-gpt2-124m-fineweb10b-2024-10-13-upstream-tpu-v2') +BENCHMARK = ('stock-gpt2-124m-fineweb10b-sgdm-layerlr-v1' if PROFILE == SGD_SPEEDRUN_PROFILE else 'llmc-gpt2-124m-fineweb10b-2024-10-13-upstream-tpu-v2') SOURCE_COMMIT = '7ecd8906afe6ed7a2b2cdb731c042f26d525b820' TOTAL_STEPS = 19560 BATCH_TOKENS = 524288 VAL_TOKENS = 10485760 WARMUP = 700 WARMDOWN = 0 -LR_SCHEDULE = 'cosine_to_zero' +LR_SCHEDULE = 'cosine_to_floor_0.10' PEAK_ADAM_LR = 0.0006 # Published GPT-2 Small SGDM center from tml-epfl/why-weight-decay. SGD_PEAK_LR = 0.3 @@ -25,6 +25,11 @@ SGD_NESTEROV = False GRAD_CLIP = 1.0 MEASUREMENT_INTERVAL = 250 +SGD_COSINE_FLOOR = 0.10 +SGD_EMBEDDING_MULTIPLIER = 0.60 +SGD_OTHER_MULTIPLIER = 0.75 +SGD_BLOCK_MULTIPLIERS = (0.75, 0.80, 0.85, 0.90, 1.20, 1.35, 1.45, 1.40, 1.30, 1.15, 1.00, 0.90) +RESUMABLE_CHECKPOINT_INTERVAL = 2000 def lr_factor(step): @@ -33,7 +38,8 @@ def lr_factor(step): if step < WARMUP: return (step + 1) / WARMUP ratio = min(1.0, (step - WARMUP) / (TOTAL_STEPS - WARMUP)) - return 0.5 * (1.0 + math.cos(math.pi * ratio)) + cosine = 0.5 * (1.0 + math.cos(math.pi * ratio)) + return SGD_COSINE_FLOOR + (1.0 - SGD_COSINE_FLOOR) * cosine def protocol(): From c68894652dba159e0bce4a3632ff724990969044 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Sat, 17 Oct 2026 17:01:09 -0700 Subject: [PATCH 2/5] Add blockwise learning rates to SGDM --- baseline/gpt2_small/muon_speedrun/optim.py | 32 ++++++++++++++++++---- 1 file changed, 26 insertions(+), 6 deletions(-) diff --git a/baseline/gpt2_small/muon_speedrun/optim.py b/baseline/gpt2_small/muon_speedrun/optim.py index 7749ca7..c97af7c 100644 --- a/baseline/gpt2_small/muon_speedrun/optim.py +++ b/baseline/gpt2_small/muon_speedrun/optim.py @@ -4,7 +4,8 @@ import torch from stock_config import ARCHITECTURE from benchmark_config import (BENCHMARK, PROFILE, SGD_SPEEDRUN_PROFILE, lr_factor, GRAD_CLIP, - SGD_PEAK_LR, SGD_MOMENTUM, SGD_WEIGHT_DECAY, SGD_NESTEROV) + SGD_PEAK_LR, SGD_MOMENTUM, SGD_WEIGHT_DECAY, SGD_NESTEROV, + SGD_EMBEDDING_MULTIPLIER, SGD_OTHER_MULTIPLIER, SGD_BLOCK_MULTIPLIERS) def schedule(step): @@ -123,10 +124,29 @@ def make_stock_optimizers(model, rt, kind): if kind == 'sgd_momentum': 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=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)] + grouped = defaultdict(list) + seen = set() + for name, p in model.named_parameters(): + if id(p) in seen: + continue + seen.add(id(p)) + if name.startswith('transformer.h.'): + block = int(name.split('.')[2]) + role = f'block_{block}' + multiplier = SGD_BLOCK_MULTIPLIERS[block] + elif name in ('transformer.wte.weight', 'transformer.wpe.weight', 'lm_head.weight'): + role = 'embedding' + multiplier = SGD_EMBEDDING_MULTIPLIER + else: + role = 'other' + multiplier = SGD_OTHER_MULTIPLIER + decay = SGD_WEIGHT_DECAY if p.ndim >= 2 else 0.0 + grouped[(role, multiplier, decay)].append(p) + groups = [] + for (role, multiplier, decay), params in grouped.items(): + peak = SGD_PEAK_LR * multiplier + groups.append(dict(params=params, role=role, lr=peak, peak_lr=peak, + lr_multiplier=multiplier, weight_decay=decay)) sgd = torch.optim.SGD(groups, momentum=SGD_MOMENTUM, nesterov=SGD_NESTEROV, dampening=0.0, foreach=False) sgd.benchmark = BENCHMARK @@ -152,7 +172,7 @@ 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'], + {'role':g['role'], 'peak_lr':g['peak_lr'], 'lr_multiplier':g.get('lr_multiplier', 1.0), 'weight_decay':g['weight_decay'], '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__, From 23c59ccdb1026dd18cd16d80600b110dc4257868 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Sat, 17 Oct 2026 17:01:14 -0700 Subject: [PATCH 3/5] Log layer LRs and retain resumable checkpoints --- baseline/gpt2_small/muon_speedrun/run.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/baseline/gpt2_small/muon_speedrun/run.py b/baseline/gpt2_small/muon_speedrun/run.py index 8405f3d..c1bef65 100644 --- a/baseline/gpt2_small/muon_speedrun/run.py +++ b/baseline/gpt2_small/muon_speedrun/run.py @@ -21,7 +21,7 @@ from benchmark_config import (BENCHMARK, TOTAL_STEPS, BATCH_TOKENS, VAL_TOKENS, MEASUREMENT_INTERVAL, WARMUP, WARMDOWN, LR_SCHEDULE, GRAD_CLIP, SGD_PEAK_LR, SGD_MOMENTUM, SGD_WEIGHT_DECAY, SGD_NESTEROV, - lr_factor, protocol) + lr_factor, protocol, RESUMABLE_CHECKPOINT_INTERVAL) TARGET = 3.28 HERE = Path(__file__).resolve().parent @@ -75,6 +75,10 @@ def save_checkpoint(root, model, muon, adam, stream, step, manifest, rt, validat temp.unlink(missing_ok=True) os.link(checkpoint, temp) temp.replace(root/name) + if step > 0 and step % RESUMABLE_CHECKPOINT_INTERVAL == 0: + archive = root/f'checkpoint_step_{step:05d}.pt' + archive.unlink(missing_ok=True) + os.link(checkpoint, archive) write_json(root/'checkpoint_latest.json', {'file':checkpoint.name, 'step':step, 'target_met':passed, 'validation':validation}) if validation is not None: @@ -274,7 +278,8 @@ def train(a): row = {'kind':'train', 'step':step, 'tokens_seen':step*BATCH_TOKENS, 'train_nll':loss_value, 'seconds':seconds, 'tokens_per_second':BATCH_TOKENS/seconds, 'elapsed_seconds':time.time()-started, 'lr_factor':lr_factor(step-1), - 'gradient_norm_before_clip':float(grad_norm.cpu())} + 'gradient_norm_before_clip':float(grad_norm.cpu()), + 'learning_rates':{g['role']:float(g['peak_lr'] * lr_factor(step-1)) for g in adam.param_groups}} if a.optimizer == 'muon_clip': row['muonclip'] = {k:float(v.cpu()) for k,v in muon.last_diagnostics.items()} if step <= 5 or step % 10 == 0: From 515b8b7e5405da88d76b314efda0918d107c41e1 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Sat, 17 Oct 2026 17:01:19 -0700 Subject: [PATCH 4/5] Record layerwise LR protocol --- baseline/gpt2_small/muon_speedrun/benchmark_config.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/baseline/gpt2_small/muon_speedrun/benchmark_config.py b/baseline/gpt2_small/muon_speedrun/benchmark_config.py index cb8ed3c..7f62454 100644 --- a/baseline/gpt2_small/muon_speedrun/benchmark_config.py +++ b/baseline/gpt2_small/muon_speedrun/benchmark_config.py @@ -49,7 +49,13 @@ def protocol(): 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=SGD_WEIGHT_DECAY, sgd_nesterov=SGD_NESTEROV, adam_betas=[0.9, 0.95], + sgd_weight_decay=SGD_WEIGHT_DECAY, sgd_nesterov=SGD_NESTEROV, + sgd_embedding_multiplier=SGD_EMBEDDING_MULTIPLIER, + sgd_other_multiplier=SGD_OTHER_MULTIPLIER, + sgd_block_multipliers=list(SGD_BLOCK_MULTIPLIERS), + sgd_cosine_floor=SGD_COSINE_FLOOR, + resumable_checkpoint_interval=RESUMABLE_CHECKPOINT_INTERVAL, + 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 ae63a8e5ce3252f17b089c90567f96ddcd13dc2a Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Sat, 17 Oct 2026 17:01:28 -0700 Subject: [PATCH 5/5] Update SGDM tests for layerwise schedule --- .../muon_speedrun/test_sgd_momentum.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/baseline/gpt2_small/muon_speedrun/test_sgd_momentum.py b/baseline/gpt2_small/muon_speedrun/test_sgd_momentum.py index 8c0ff3b..5772699 100644 --- a/baseline/gpt2_small/muon_speedrun/test_sgd_momentum.py +++ b/baseline/gpt2_small/muon_speedrun/test_sgd_momentum.py @@ -11,7 +11,8 @@ 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, SGD_WEIGHT_DECAY, TOTAL_STEPS, WARMUP, WARMDOWN, lr_factor +from benchmark_config import (SGD_PEAK_LR, SGD_MOMENTUM, SGD_WEIGHT_DECAY, TOTAL_STEPS, WARMUP, + SGD_OTHER_MULTIPLIER, SGD_BLOCK_MULTIPLIERS, SGD_COSINE_FLOOR, lr_factor) class TinyTiedModel(torch.nn.Module): @@ -37,8 +38,9 @@ def test_tied_parameter_once_and_classical_momentum(self): self.assertEqual(len(params), len(set(map(id, params)))) self.assertEqual(set(map(id, params)), set(map(id, model.parameters()))) 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.assertAlmostEqual(groups['block_0']['lr_multiplier'], SGD_BLOCK_MULTIPLIERS[0]) + self.assertEqual(groups['block_0']['weight_decay'], SGD_WEIGHT_DECAY) + self.assertEqual(groups['other']['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") @@ -47,18 +49,17 @@ def test_tied_parameter_once_and_classical_momentum(self): 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))) + self.assertTrue(torch.allclose(p, old - SGD_PEAK_LR * SGD_OTHER_MULTIPLIER * 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))) + self.assertTrue(torch.allclose(p, previous - SGD_PEAK_LR * SGD_OTHER_MULTIPLIER * lr_factor(1) * (1 + SGD_MOMENTUM))) def test_speedrun_schedule(self): - self.assertEqual(TOTAL_STEPS, 6200) + self.assertEqual(TOTAL_STEPS, 19560) 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) + self.assertAlmostEqual(lr_factor(TOTAL_STEPS), SGD_COSINE_FLOOR) if __name__ == "__main__":