fix(train): fix regression from #971 - keep Muon-excluded matrices at Muon's LR

src/speculators/train/optimizers.py

@@ -3,9 +3,10 @@
 Provides a single entry point, :func:`build_optimizers`, that returns the list of
 optimizers the trainer should drive. The default ("adamw") returns a single AdamW
 optimizer over all parameters, preserving the historical behavior. The "muon" option
-returns two optimizers: ``torch.optim.Muon`` over the 2D weight matrices (which is all
-Muon supports) and ``torch.optim.AdamW`` over everything else (norms, biases, and the
-embedding / LM-head matrices, following standard Muon practice).
+returns two optimizers: ``torch.optim.Muon`` over the supported 2D weight matrices and
+``torch.optim.AdamW`` over everything else. Fresh, output-adjacent token-transition
+factors use a dedicated AdamW parameter group at ``muon_lr``; other Muon exclusions
+retain the base AdamW hyperparameters.
 
 Muon works transparently for both single-GPU and multi-GPU (FSDP2) training: when the
 model is sharded with ``fully_shard`` the parameters become ``DTensor``s and Muon's
@@ -16,54 +17,81 @@
 
 import torch
 from torch import Tensor
-from torch.nn import Module
+from torch.nn import Embedding, Module
 
 logger = logging.getLogger("speculators")
 
-# Names of parameters that are 2D but should still be optimized with AdamW rather than
-# Muon, following the convention from Keller Jordan's Muon (embeddings, embedding-like
-# codebooks, Markov vocabulary factors, and output heads are excluded from the
-# orthogonalized update).
+# Muon is intended for feature-to-feature matrices, not vocabulary-indexed tables.
+# Embeddings are also detected structurally; these hints cover output heads and custom
+# codebooks implemented as raw parameters.
 _ADAMW_NAME_HINTS = (
     "embed_tokens",
     "lm_head",
     "codebook",
-    "markov_w1",
-    "markov_w2",
+)
+
+# Fresh, output-adjacent token-transition factors use AdamW for their vocabulary axes
+# but keep Muon's larger learning rate and weight decay. Match complete suffixes so
+# unrelated embeddings and codebooks remain in the base AdamW group.
+_TRANSITION_PARAM_SUFFIXES = (
+    "markov_w1.weight",
+    "markov_w2.weight",
+    "predecessor_codebook",
+    "predecessor_codebook.weight",
+    "successor_codebook",
+    "successor_codebook.weight",
 )
 
 # Muon only orthogonalizes 2D weight matrices.
 _MATRIX_NDIM = 2
 
 
+def _matches_param_suffix(name: str, suffixes: tuple[str, ...]) -> bool:
+    return any(name == suffix or name.endswith(f".{suffix}") for suffix in suffixes)
+
+
 def split_named_params_for_muon(
     model: Module,
-) -> tuple[list[tuple[str, Tensor]], list[tuple[str, Tensor]]]:
-    """Split a model's trainable parameters into Muon and AdamW groups.
-
-    A parameter goes to Muon iff it requires gradients, is a 2D matrix with both
-    dimensions > 1, and is not an embedding, codebook, or vocabulary-output weight;
-    everything else goes to AdamW. Degenerate 2D weights (``[1, N]`` / ``[N, 1]``
-    vectors) route to AdamW -- Muon orthogonalizes matrices, not vectors, and crashes
-    on them under FSDP2.
+) -> tuple[
+    list[tuple[str, Tensor]], list[tuple[str, Tensor]], list[tuple[str, Tensor]]
+]:
+    """Split trainable parameters by optimizer and AdamW hyperparameters.
+
+    Fresh token-transition factors go to ``transition_params`` so AdamW can update
+    them with Muon's learning rate and weight decay. Other embeddings, codebooks, and
+    vocabulary-output weights use the base AdamW group. A remaining parameter goes to
+    Muon iff it is a 2D matrix with both dimensions > 1; norms, biases, and degenerate
+    2D weights (``[1, N]`` / ``[N, 1]`` vectors) use base AdamW because Muon
+    orthogonalizes matrices, not vectors, and crashes on them under FSDP2.
 
     :param model: The model whose parameters should be partitioned.
-    :return: A ``(muon_params, adamw_params)`` tuple of named parameter lists.
+    :return: A ``(muon_params, adamw_params, transition_params)`` tuple.
     """
+    embedding_param_ids = {
+        id(param)
+        for module in model.modules()
+        if isinstance(module, Embedding)
+        for param in module.parameters(recurse=False)
+    }
+
     muon_params: list[tuple[str, Tensor]] = []
     adamw_params: list[tuple[str, Tensor]] = []
+    transition_params: list[tuple[str, Tensor]] = []
     for name, param in model.named_parameters():
         if not param.requires_grad:
             continue
-        if (
-            param.ndim == _MATRIX_NDIM
-            and min(param.shape) > 1  # exclude degenerate [1, N] / [N, 1] vectors
-            and not any(hint in name for hint in _ADAMW_NAME_HINTS)
+        if _matches_param_suffix(name, _TRANSITION_PARAM_SUFFIXES):
+            transition_params.append((name, param))
+        elif (
+            param.ndim != _MATRIX_NDIM
+            or min(param.shape) == 1
+            or id(param) in embedding_param_ids
+            or any(hint in name for hint in _ADAMW_NAME_HINTS)
         ):
-            muon_params.append((name, param))
-        else:
             adamw_params.append((name, param))
-    return muon_params, adamw_params
+        else:
+            muon_params.append((name, param))
+    return muon_params, adamw_params, transition_params
 
 
 def build_optimizers(model: Module, config) -> list[torch.optim.Optimizer]:
@@ -72,7 +100,8 @@ def build_optimizers(model: Module, config) -> list[torch.optim.Optimizer]:
     :param model: The model to optimize.
     :param config: A ``TrainerConfig`` holding the optimizer hyperparameters.
     :return: A list of optimizers for the trainer to step in tandem. The default
-        "adamw" returns a single optimizer; "muon" returns ``[Muon, AdamW]``.
+        "adamw" returns a single optimizer; "muon" returns ``[Muon, AdamW]`` when
+        both parameter types are present.
     """
     if config.optimizer == "adamw":
         return [
@@ -84,11 +113,15 @@ def build_optimizers(model: Module, config) -> list[torch.optim.Optimizer]:
         ]
 
     if config.optimizer == "muon":
-        muon_params, adamw_params = split_named_params_for_muon(model)
+        muon_params, adamw_params, transition_params = split_named_params_for_muon(
+            model
+        )
         logger.info(
-            "Muon optimizer: %d 2D params via Muon, %d params via AdamW.",
+            "Muon optimizer: %d via Muon, %d via base AdamW, %d transition factors "
+            "via AdamW at muon_lr.",
             len(muon_params),
             len(adamw_params),
+            len(transition_params),
         )
 
         optimizers: list[torch.optim.Optimizer] = []
@@ -103,10 +136,21 @@ def build_optimizers(model: Module, config) -> list[torch.optim.Optimizer]:
                     adjust_lr_fn=config.muon_adjust_lr_fn,
                 )
             )
+        adamw_param_groups = []
         if adamw_params:
+            adamw_param_groups.append({"params": adamw_params})
+        if transition_params:
+            adamw_param_groups.append(
+                {
+                    "params": transition_params,
+                    "lr": config.muon_lr,
+                    "weight_decay": config.muon_weight_decay,
+                }
+            )
+        if adamw_param_groups:
             optimizers.append(
                 torch.optim.AdamW(
-                    adamw_params,
+                    adamw_param_groups,
                     lr=config.lr,
                     weight_decay=config.weight_decay,
                 )

tests/unit/models/test_dflash2_model_definitions.py

@@ -225,12 +225,11 @@ def test_candidate_codebooks_use_adamw_under_muon_optimizer():
         top_k=3,
     )
 
-    muon, adamw = split_named_params_for_muon(selector)
-    muon_names = {name for name, _ in muon}
-    adamw_names = {name for name, _ in adamw}
+    muon, adamw, transition = split_named_params_for_muon(selector)
 
-    assert muon_names == {"hidden_projection.weight"}
-    assert adamw_names == {
+    assert {name for name, _ in muon} == {"hidden_projection.weight"}
+    assert not adamw
+    assert {name for name, _ in transition} == {
         "predecessor_codebook",
         "successor_codebook",
     }

tests/unit/models/test_dspark_model_definitions.py

@@ -1,10 +1,15 @@
 """Unit tests for DSpark Markov and confidence heads."""
 
+from types import SimpleNamespace
+
 import pytest
 import torch
 
 from speculators.models.dspark.model_definitions import ConfidenceHead, MarkovHead
-from speculators.train.optimizers import split_named_params_for_muon
+from speculators.train.optimizers import (
+    build_optimizers,
+    split_named_params_for_muon,
+)
 
 
 class TestMarkovHead:
@@ -55,17 +60,49 @@ def test_lookup_embedding_has_small_initialization(self):
     def test_vocab_factors_use_adamw_under_muon_optimizer(self):
         head = self._head("gated")
 
-        muon, adamw = split_named_params_for_muon(head)
-        muon_names = {name for name, _ in muon}
-        adamw_names = {name for name, _ in adamw}
+        muon, adamw, transition = split_named_params_for_muon(head)
 
-        assert muon_names == {"gate_proj.weight"}
-        assert adamw_names == {
+        assert {name for name, _ in muon} == {"gate_proj.weight"}
+        assert {name for name, _ in adamw} == {"gate_proj.bias"}
+        assert {name for name, _ in transition} == {
             "markov_w1.weight",
             "markov_w2.weight",
-            "gate_proj.bias",
         }
 
+    def test_vocab_factors_keep_muon_lr_not_the_base_lr(self):
+        """The Markov factors are skipped by Muon because a vocabulary index is not a
+        feature axis -- not because they want a 10x smaller step than their neighbours.
+        """
+        head = self._head("gated")
+        config = SimpleNamespace(
+            optimizer="muon",
+            lr=3e-4,
+            weight_decay=0.01,
+            muon_lr=3e-3,
+            muon_momentum=0.95,
+            muon_weight_decay=0.1,
+            muon_ns_steps=5,
+            muon_adjust_lr_fn="match_rms_adamw",
+        )
+
+        muon, adamw = build_optimizers(head, config)
+
+        assert isinstance(muon, torch.optim.Muon)
+        assert isinstance(adamw, torch.optim.AdamW)
+        assert len(adamw.param_groups) == 2
+
+        groups = {
+            name: group
+            for group in adamw.param_groups
+            for name in group.get("param_names") or []
+        }
+
+        assert groups["gate_proj.bias"]["lr"] == config.lr
+        assert groups["gate_proj.bias"]["weight_decay"] == config.weight_decay
+        for name in ("markov_w1.weight", "markov_w2.weight"):
+            assert groups[name]["lr"] == config.muon_lr
+            assert groups[name]["weight_decay"] == config.muon_weight_decay
+
     def test_invalid_rank_raises(self):
         with pytest.raises(ValueError):
             MarkovHead(

tests/unit/train/test_optimizers.py

@@ -0,0 +1,24 @@
+"""Tests for optimizer parameter partitioning."""
+
+import torch
+from torch import nn
+
+from speculators.train.optimizers import split_named_params_for_muon
+
+
+def test_generic_embeddings_heads_and_codebooks_use_base_adamw():
+    model = nn.Module()
+    model.token_lookup = nn.Embedding(32, 8)
+    model.lm_head = nn.Linear(8, 32, bias=False)
+    model.aux_codebook = nn.Parameter(torch.empty(32, 8))
+    model.projection = nn.Linear(8, 16, bias=False)
+
+    muon, adamw, transition = split_named_params_for_muon(model)
+
+    assert {name for name, _ in muon} == {"projection.weight"}
+    assert {name for name, _ in adamw} == {
+        "token_lookup.weight",
+        "lm_head.weight",
+        "aux_codebook",
+    }
+    assert not transition