| """Muon optimizer + Muon/AdamW hybrid.""" |
|
|
| import pytest |
| import torch |
|
|
| from matilda import Transformer, ModelConfig |
| from matilda.optim import ( |
| build_optimizer, cosine_warmup_scheduler, build_scheduler, |
| wsd_scheduler, WarmupStableDecay, Muon, HybridOptimizer, |
| zeropower_via_newtonschulz5, |
| ) |
|
|
|
|
| def test_newtonschulz_orthogonalizes(): |
| torch.manual_seed(0) |
| G = torch.randn(32, 16) |
| X = zeropower_via_newtonschulz5(G, steps=5).float() |
| |
| s = torch.linalg.svdvals(X) |
| assert (s > 0.5).all() and (s < 1.5).all() |
|
|
|
|
| def test_hybrid_splits_params(): |
| cfg = ModelConfig(vocab_size=128, max_seq_len=32, d_model=64, |
| n_layers=2, n_heads=4, n_kv_heads=2) |
| opt = build_optimizer(Transformer(cfg), name="muon") |
| assert isinstance(opt, HybridOptimizer) |
| muon, adamw = opt.optimizers |
| assert isinstance(muon, Muon) |
| |
| assert all(p.ndim == 2 for g in muon.param_groups for p in g["params"]) |
|
|
|
|
| @pytest.mark.slow |
| def test_muon_overfits_single_batch(): |
| cfg = ModelConfig(vocab_size=256, max_seq_len=64, d_model=128, |
| n_layers=2, n_heads=4, n_kv_heads=2) |
| model = Transformer(cfg).train() |
| torch.manual_seed(0) |
| idx = torch.randint(0, cfg.vocab_size, (4, 32)) |
| tgt = torch.randint(0, cfg.vocab_size, (4, 32)) |
| opt = build_optimizer(model, name="muon", lr=3e-3, muon_lr=0.02) |
| sched = cosine_warmup_scheduler(opt, warmup_steps=10, total_steps=300) |
| last = None |
| for _ in range(300): |
| _, loss = model(idx, tgt) |
| opt.zero_grad(set_to_none=True) |
| loss.backward() |
| opt.step() |
| sched.step() |
| last = loss.item() |
| assert last < 0.5, f"Muon failed to overfit; final loss={last:.3f}" |
|
|
|
|
| def _dummy_opt(base_lr=1.0): |
| """Minimal optimizer just to exercise scheduler interface.""" |
| p = torch.nn.Parameter(torch.zeros(2, 2)) |
| return torch.optim.SGD([p], lr=base_lr) |
|
|
|
|
| def test_wsd_schedule_three_phases(): |
| """Warmup ramps linearly to peak, stable holds peak, decay drops to min.""" |
| base = 1.0 |
| warmup, total, stable = 10, 100, 0.8 |
| sched = wsd_scheduler(_dummy_opt(base), warmup_steps=warmup, |
| total_steps=total, stable_share=stable, |
| min_lr_ratio=0.1) |
| lrs = [] |
| for _ in range(total): |
| lrs.append(sched.get_last_lr()[0]) |
| sched.step() |
| |
| assert lrs[0] == pytest.approx(base / warmup, rel=1e-6) |
| assert lrs[warmup - 1] == pytest.approx(base, rel=1e-6) |
| |
| decay_start = warmup + int(stable * (total - warmup)) |
| assert lrs[warmup] == pytest.approx(base, rel=1e-6) |
| assert lrs[decay_start - 1] == pytest.approx(base, rel=1e-6) |
| |
| decay_lrs = lrs[decay_start:] |
| assert all(decay_lrs[i] >= decay_lrs[i + 1] for i in range(len(decay_lrs) - 1)) |
| |
| assert lrs[-1] < base |
| assert lrs[-1] >= 0.1 * base * 0.5 |
|
|
|
|
| def test_wsd_state_dict_roundtrip(): |
| """Resume must restore step + base_lrs and re-apply the LR exactly.""" |
| opt = _dummy_opt(1.0) |
| s1 = wsd_scheduler(opt, warmup_steps=10, total_steps=100, stable_share=0.8) |
| for _ in range(50): |
| s1.step() |
| sd = s1.state_dict() |
| s2 = wsd_scheduler(_dummy_opt(1.0), warmup_steps=10, total_steps=100, |
| stable_share=0.8) |
| s2.load_state_dict(sd) |
| assert s2.last_step == s1.last_step |
| assert s2.get_last_lr()[0] == pytest.approx(s1.get_last_lr()[0], rel=1e-9) |
|
|
|
|
| def test_build_scheduler_dispatch(): |
| assert isinstance( |
| build_scheduler("wsd", _dummy_opt(), 10, 100), WarmupStableDecay) |
| cos = build_scheduler("cosine", _dummy_opt(), 10, 100) |
| assert hasattr(cos, "step") and hasattr(cos, "get_last_lr") |
| with pytest.raises(ValueError): |
| build_scheduler("bogus", _dummy_opt(), 10, 100) |
|
|
|
|
| def test_hybrid_state_dict_roundtrip(): |
| cfg = ModelConfig(vocab_size=128, max_seq_len=32, d_model=64, |
| n_layers=2, n_heads=4, n_kv_heads=2) |
| model = Transformer(cfg) |
| opt = build_optimizer(model, name="muon") |
| |
| _, loss = model(torch.randint(0, 128, (2, 16)), torch.randint(0, 128, (2, 16))) |
| loss.backward() |
| opt.step() |
| sd = opt.state_dict() |
| opt2 = build_optimizer(Transformer(cfg), name="muon") |
| opt2.load_state_dict(sd) |
| assert len(sd["opts"]) == 2 |
|
|
