upload muon.py

muon.py

@@ -0,0 +1,121 @@
+"""
+Muon Optimizer — Keller Jordan, NanoGPT speedrun.
+
+Newton-Schulz iterasyonu ile orthogonalize edilmiş momentum.
+2D ağırlıklara (Linear weight) uygulanır. 1D parametreler (norm weight,
+bias, embedding) AdamW'ye verilir.
+
+Referans:
+  https://github.com/KellerJordan/modded-nanogpt
+  https://kellerjordan.github.io/posts/muon/
+
+Kullanim:
+    # Param ayri:
+    muon_params = [p for p in model.parameters() if p.dim() >= 2 and p.requires_grad]
+    other_params = [p for p in model.parameters() if p.dim() < 2 and p.requires_grad]
+
+    # Embedding ve lm_head'i muon'dan ayir (yaygin best practice)
+    embed_params = [model.wte.weight]  # tied ise lm_head dahil
+    muon_params = [p for p in muon_params if not any(p is e for e in embed_params)]
+    other_params = other_params + embed_params
+
+    optimizer_muon = Muon(muon_params, lr=2e-2, momentum=0.95)
+    optimizer_adam = torch.optim.AdamW(other_params, lr=3e-4, ...)
+"""
+
+import torch
+
+
+@torch.no_grad()
+def newton_schulz(G: torch.Tensor, steps: int = 5) -> torch.Tensor:
+    """G matrisini orthogonalize et (yaklasik USV^T -> UV^T).
+
+    Newton-Schulz quintic iteration. bf16'da kararli, hizli.
+    """
+    assert G.ndim == 2
+    a, b, c = (3.4445, -4.7750, 2.0315)
+    X = G.to(torch.bfloat16)
+    # Boyut yonune gore transpose (her iki yonde de calissin)
+    if X.size(0) > X.size(1):
+        X = X.T
+    # Spektral normu yaklasik 1'e cek
+    X = X / (X.norm() + 1e-7)
+    for _ in range(steps):
+        A = X @ X.T
+        B = b * A + c * (A @ A)
+        X = a * X + B @ X
+    if G.size(0) > G.size(1):
+        X = X.T
+    return X.to(G.dtype)
+
+
+class Muon(torch.optim.Optimizer):
+    """Muon: Momentum + orthogonalize edilmiş update.
+
+    Sadece 2D parametreler için. 1D'leri AdamW ile ayrı eğit.
+
+    Args:
+        params: 2D parametreler iterable
+        lr: 0.02 (AdamW'nin ~50x'i, çünkü update'ler ortonormal)
+        momentum: 0.95
+        nesterov: True (genelde daha iyi)
+        ns_steps: Newton-Schulz iter sayisi (5 default)
+    """
+    def __init__(self, params, lr=0.02, momentum=0.95, nesterov=True, ns_steps=5):
+        defaults = dict(lr=lr, momentum=momentum, nesterov=nesterov, ns_steps=ns_steps)
+        super().__init__(params, defaults)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            lr = group["lr"]
+            momentum = group["momentum"]
+            nesterov = group["nesterov"]
+            ns_steps = group["ns_steps"]
+
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                if p.ndim < 2:
+                    raise ValueError(
+                        f"Muon sadece >=2D param destekler, {p.ndim}D bulundu. "
+                        "1D paramları AdamW'ye ver.")
+
+                g = p.grad
+                state = self.state[p]
+                if "momentum_buffer" not in state:
+                    state["momentum_buffer"] = torch.zeros_like(g)
+
+                buf = state["momentum_buffer"]
+                buf.mul_(momentum).add_(g)
+
+                # Nesterov momentum
+                if nesterov:
+                    g = g.add(buf, alpha=momentum)
+                else:
+                    g = buf
+
+                # Reshape if needed (e.g., conv weight)
+                original_shape = g.shape
+                if g.ndim > 2:
+                    g = g.view(g.size(0), -1)
+
+                # Newton-Schulz orthogonalization
+                g_orth = newton_schulz(g, steps=ns_steps)
+
+                # Scale: sqrt(max(out, in) / min(out, in)) — ~spectral norm
+                # Ya da basitce sqrt(d_out / d_in) gibi.
+                # Modded-nanogpt: scale = max(1, p.shape[0]/p.shape[1]) ** 0.5
+                scale = max(1.0, g_orth.size(0) / g_orth.size(1)) ** 0.5
+
+                # Geri reshape
+                g_orth = g_orth.view(original_shape)
+
+                p.add_(g_orth, alpha=-lr * scale)
+
+        return loss