chesstransformer/optimizer.py

"""Combined AdamW + Muon optimizer.

The two underlying optimizers are kept as separate ``torch.optim`` instances
internally; this class is a thin wrapper that exposes the standard
``torch.optim.Optimizer`` interface (``step``, ``zero_grad``, ``param_groups``,
``state_dict`` / ``load_state_dict``) so that LR schedulers, gradient
accumulation, and HuggingFace ``Accelerator`` treat it as a single optimizer.

Routing rule (when constructed from an ``nn.Module``):

* parameters whose name contains an ``embedding_keyword`` → AdamW (low LR)
* parameters whose name contains a ``head_keyword``      → AdamW (no decay)
* remaining parameters with ``ndim >= 2``                → Muon
* remaining parameters with ``ndim < 2`` (biases/norms)  → AdamW (no decay)
"""

from __future__ import annotations

from typing import Iterable

import torch
from torch import nn
from torch.optim import AdamW, Muon, Optimizer


class AdamWMuon(Optimizer):
    """Wrapper that drives an inner ``AdamW`` and ``Muon`` in lock-step.

    Parameters
    ----------
    model:
        ``nn.Module`` whose parameters will be auto-routed by name.
    lr:
        Base AdamW learning rate (used for the ``other`` 1D parameter group).
    lr_embedding, lr_head:
        Per-group AdamW LR overrides. Default to ``lr``.
    lr_muon:
        Muon LR for ≥2D body weights. Default ``= 3 * lr`` (Muon tolerates a
        higher LR than AdamW because updates are orthogonalised).
    weight_decay:
        AdamW weight decay applied to the embedding group. Heads and 1D
        params get ``0.0`` by convention.
    weight_decay_muon:
        Muon weight decay. Default = ``weight_decay``.
    embedding_keywords / head_keywords:
        Substrings used to classify parameters by ``named_parameters`` key.
    """

    def __init__(
        self,
        model: nn.Module,
        *,
        lr: float = 3e-4,
        lr_embedding: float | None = None,
        lr_head: float | None = None,
        lr_muon: float | None = None,
        weight_decay: float = 0.1,
        weight_decay_muon: float | None = None,
        betas: tuple[float, float] = (0.9, 0.95),
        eps: float = 1e-8,
        muon_momentum: float = 0.95,
        muon_nesterov: bool = True,
        muon_ns_steps: int = 5,
        embedding_keywords: Iterable[str] = ("embedding", "cls_token"),
        head_keywords: Iterable[str] = ("lm_head", "value_head"),
    ) -> None:
        if not isinstance(model, nn.Module):
            raise TypeError(
                f"AdamWMuon expects an nn.Module to auto-route parameters by name, "
                f"got {type(model).__name__}."
            )

        lr_embedding = lr if lr_embedding is None else lr_embedding
        lr_head = lr if lr_head is None else lr_head
        lr_muon = 3.0 * lr if lr_muon is None else lr_muon
        weight_decay_muon = weight_decay if weight_decay_muon is None else weight_decay_muon

        embedding_keywords = tuple(embedding_keywords)
        head_keywords = tuple(head_keywords)

        emb_params: list[torch.Tensor] = []
        head_params: list[torch.Tensor] = []
        other_params: list[torch.Tensor] = []
        muon_params: list[torch.Tensor] = []

        for name, param in model.named_parameters():
            if not param.requires_grad:
                continue
            if any(kw in name for kw in embedding_keywords):
                emb_params.append(param)
            elif any(kw in name for kw in head_keywords):
                head_params.append(param)
            elif param.ndim >= 2:
                muon_params.append(param)
            else:
                other_params.append(param)

        adamw_groups: list[dict] = []
        if emb_params:
            adamw_groups.append({
                "params": emb_params, "lr": lr_embedding,
                "weight_decay": weight_decay, "name": "adamw_embedding",
            })
        if other_params:
            adamw_groups.append({
                "params": other_params, "lr": lr,
                "weight_decay": 0.0, "name": "adamw_other",
            })
        if head_params:
            adamw_groups.append({
                "params": head_params, "lr": lr_head,
                "weight_decay": 0.0, "name": "adamw_head",
            })

        if not adamw_groups:
            raise ValueError("No AdamW-eligible parameters found in model.")

        self.adamw = AdamW(adamw_groups, lr=lr, betas=betas, eps=eps,
                           weight_decay=weight_decay)

        if muon_params:
            self.muon: Muon | None = Muon(
                [{"params": muon_params, "lr": lr_muon,
                  "weight_decay": weight_decay_muon, "name": "muon_body"}],
                lr=lr_muon, weight_decay=weight_decay_muon,
                momentum=muon_momentum, nesterov=muon_nesterov,
                ns_steps=muon_ns_steps,
            )
        else:
            self.muon = None

        # Initialise the ``Optimizer`` base with a placeholder group so
        # ``isinstance(opt, Optimizer)`` checks (Accelerate, schedulers) pass
        # and ``self.state`` / ``self.defaults`` exist. We then expose the
        # *inner* groups via ``param_groups`` so schedulers mutate them in
        # place and both sub-optimizers see the new LR.
        all_params = emb_params + other_params + head_params + muon_params
        self._init_done = False
        super().__init__([{"params": all_params}], defaults={"lr": lr})
        self.param_groups = self._combined_param_groups()
        self._init_done = True

    # ── helpers ─────────────────────────────────────────────────────────
    def _combined_param_groups(self) -> list[dict]:
        groups = list(self.adamw.param_groups)
        if self.muon is not None:
            groups += list(self.muon.param_groups)
        return groups

    # ── Optimizer API ───────────────────────────────────────────────────
    @torch.no_grad()
    def step(self, closure=None):
        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()
        self.adamw.step()
        if self.muon is not None:
            self.muon.step()
        return loss

    def zero_grad(self, set_to_none: bool = True) -> None:
        self.adamw.zero_grad(set_to_none=set_to_none)
        if self.muon is not None:
            self.muon.zero_grad(set_to_none=set_to_none)

    def state_dict(self) -> dict:
        return {
            "adamw": self.adamw.state_dict(),
            "muon": self.muon.state_dict() if self.muon is not None else None,
        }

    def load_state_dict(self, state_dict: dict) -> None:
        self.adamw.load_state_dict(state_dict["adamw"])
        if self.muon is not None and state_dict.get("muon") is not None:
            self.muon.load_state_dict(state_dict["muon"])
        # Re-link in case sub-optimizers rebuilt their group dicts.
        self.param_groups = self._combined_param_groups()

    def add_param_group(self, param_group) -> None:
        if not getattr(self, "_init_done", False):
            # Called from ``Optimizer.__init__`` with the placeholder group.
            super().add_param_group(param_group)
            return
        raise NotImplementedError(
            "add_param_group is not supported on AdamWMuon after construction; "
            "build a new optimizer with the full parameter set instead."
        )

    def __repr__(self) -> str:
        n_emb = sum(p.numel() for g in self.adamw.param_groups
                    if g.get("name") == "adamw_embedding" for p in g["params"])
        n_head = sum(p.numel() for g in self.adamw.param_groups
                     if g.get("name") == "adamw_head" for p in g["params"])
        n_other = sum(p.numel() for g in self.adamw.param_groups
                      if g.get("name") == "adamw_other" for p in g["params"])
        n_muon = (sum(p.numel() for g in self.muon.param_groups for p in g["params"])
                  if self.muon is not None else 0)
        return (
            f"AdamWMuon(emb={n_emb:,}, other_1d={n_other:,}, "
            f"head={n_head:,}, muon_2d={n_muon:,})"
        )


# Backwards-compat alias for the previous skeleton class.
AdamWMuonOptim = AdamWMuon