Add layerwise SGDM learning rates for stock GPT-2

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():
@@ -43,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,

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__,

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:

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__":