From 171a907651916303ba34afba916a607510eaa01b Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Sat, 17 Oct 2026 17:03:08 -0700 Subject: [PATCH 1/2] Separate SGDM matrix decay groups --- baseline/gpt2_small/muon_speedrun/optim.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/baseline/gpt2_small/muon_speedrun/optim.py b/baseline/gpt2_small/muon_speedrun/optim.py index c97af7c..8610314 100644 --- a/baseline/gpt2_small/muon_speedrun/optim.py +++ b/baseline/gpt2_small/muon_speedrun/optim.py @@ -140,8 +140,11 @@ def make_stock_optimizers(model, rt, kind): else: role = 'other' multiplier = SGD_OTHER_MULTIPLIER + # Keep matrix decay and bias/LayerNorm no-decay groups separate even + # when they share a blockwise LR multiplier. decay = SGD_WEIGHT_DECAY if p.ndim >= 2 else 0.0 - grouped[(role, multiplier, decay)].append(p) + decay_role = role + ('_matrices' if p.ndim >= 2 else '_bias_and_layernorm') + grouped[(decay_role, multiplier, decay)].append(p) groups = [] for (role, multiplier, decay), params in grouped.items(): peak = SGD_PEAK_LR * multiplier From cd25d22fad21e056128e439bf1a82c1977878b6a Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Sat, 17 Oct 2026 17:03:14 -0700 Subject: [PATCH 2/2] Test separate SGDM decay groups --- baseline/gpt2_small/muon_speedrun/test_sgd_momentum.py | 7 ++++--- 1 file changed, 4 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 5772699..3547c30 100644 --- a/baseline/gpt2_small/muon_speedrun/test_sgd_momentum.py +++ b/baseline/gpt2_small/muon_speedrun/test_sgd_momentum.py @@ -38,9 +38,10 @@ 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.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.assertAlmostEqual(groups['block_0_matrices']['lr_multiplier'], SGD_BLOCK_MULTIPLIERS[0]) + self.assertEqual(groups['block_0_matrices']['weight_decay'], SGD_WEIGHT_DECAY) + self.assertEqual(groups['block_0_bias_and_layernorm']['weight_decay'], 0.0) + self.assertEqual(groups['other_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")