cortex/optim.py

"""Optimizer: Muon for 2D weight matrices, AdamW for everything else.

Routing is explicit via ``optax.multi_transform``: any 2D parameter that is *not*
the token-embedding table goes to Muon (all the Linear ``kernel`` weights);
embeddings, RMSNorm scales, LayerScale vectors, biases -> AdamW.

Both branches follow a Warmup-Stable-Decay (WSD) schedule, which fits Kaggle's
preemptible 9h sessions: hold the stable phase across many sessions, decay only at
the very end.
"""

from __future__ import annotations

from collections import defaultdict

import jax
import jax.numpy as jnp
import optax
from optax import contrib
from optax.contrib import MuonDimensionNumbers

from .config import TrainConfig


# ----------------------------------------------------------------------------
# WSD schedule
# ----------------------------------------------------------------------------
def wsd_schedule(peak_lr: float, warmup_steps: int, stable_steps: int,
                 decay_steps: int, end_lr_frac: float = 0.0,
                 decay_type: str = "1-sqrt") -> optax.Schedule:
    end_lr = peak_lr * end_lr_frac
    warmup = optax.linear_schedule(0.0, peak_lr, max(warmup_steps, 1))
    stable = optax.constant_schedule(peak_lr)

    if decay_type == "linear":
        decay = optax.linear_schedule(peak_lr, end_lr, max(decay_steps, 1))
    elif decay_type == "cosine":
        alpha = end_lr / peak_lr if peak_lr > 0 else 0.0
        decay = optax.cosine_decay_schedule(peak_lr, max(decay_steps, 1), alpha=alpha)
    elif decay_type == "1-sqrt":   # MiniCPM-style; sharp-then-flat
        d = max(decay_steps, 1)

        def decay(step):
            frac = jnp.clip(step / d, 0.0, 1.0)
            return end_lr + (peak_lr - end_lr) * (1.0 - jnp.sqrt(frac))
    else:
        raise ValueError(f"unknown decay_type: {decay_type}")

    return optax.join_schedules([warmup, stable, decay],
                                [warmup_steps, warmup_steps + stable_steps])


# ----------------------------------------------------------------------------
# param routing
# ----------------------------------------------------------------------------
def _is_muon_leaf(path, x) -> bool:
    """A param goes to Muon iff it is a 2D matrix that is not the embedding table."""
    key = jax.tree_util.keystr(path)
    return getattr(x, "ndim", 0) == 2 and "embedding" not in key


def muon_dimension_numbers(params):
    """Routing for ``optax.contrib.muon``.

    Returns a prefix tree: ``MuonDimensionNumbers()`` (reduction_axis=0,
    output_axis=1 -> correct for ``[in, out]`` Linear kernels) for params Muon
    should orthogonalize, and ``None`` for params that should fall through to the
    AdamW branch (embeddings, RMSNorm scales, LayerScale vectors, biases).
    """
    return jax.tree_util.tree_map_with_path(
        lambda p, x: MuonDimensionNumbers() if _is_muon_leaf(p, x) else None,
        params)


def param_stats(params) -> dict:
    counts = defaultdict(int)

    def f(path, x):
        lbl = "muon" if _is_muon_leaf(path, x) else "adamw"
        counts[lbl] += int(x.size)

    jax.tree_util.tree_map_with_path(f, params)
    counts["total"] = counts["muon"] + counts["adamw"]
    return dict(counts)


# ----------------------------------------------------------------------------
# build optimizer
# ----------------------------------------------------------------------------
def build_optimizer(tcfg: TrainConfig) -> optax.GradientTransformation:
    """Muon (2D matmuls) + AdamW (rest), both on a WSD schedule, with grad clipping.

    Uses the full ``optax.contrib.muon`` so we get the Moonlight update-scale
    matching (``scale_by_shape``) that lets a single sane LR work; routing of the
    embedding table to AdamW is handled by ``muon_dimension_numbers``.
    """
    muon_sched = wsd_schedule(tcfg.muon_lr, tcfg.warmup_steps, tcfg.stable_steps,
                              tcfg.decay_steps, tcfg.end_lr_frac, tcfg.decay_type)
    adam_sched = wsd_schedule(tcfg.adam_lr, tcfg.warmup_steps, tcfg.stable_steps,
                              tcfg.decay_steps, tcfg.end_lr_frac, tcfg.decay_type)

    tx_core = contrib.muon(
        learning_rate=muon_sched,
        beta=tcfg.muon_beta,
        weight_decay=tcfg.weight_decay,
        nesterov=True,
        adam_b1=tcfg.adam_b1,
        adam_b2=tcfg.adam_b2,
        adam_weight_decay=tcfg.weight_decay,
        adam_learning_rate=adam_sched,
        mu_dtype=jnp.dtype(tcfg.opt_mu_dtype),   # bf16 first moments (Muon mu + AdamW m); v stays fp32
        muon_weight_dimension_numbers=muon_dimension_numbers,
    )
    return optax.chain(optax.clip_by_global_norm(tcfg.grad_clip), tx_core)