feat(train): add Dion3 as an opt-in matrix optimizer

scripts/benchmark.py

@@ -341,6 +341,8 @@ def run_benchmark(bench_args, train_args) -> dict:
         muon_weight_decay=train_args.muon_weight_decay,
         muon_ns_steps=train_args.muon_ns_steps,
         muon_adjust_lr_fn=train_args.muon_adjust_lr_fn,
+        dion_fraction=train_args.dion_fraction,
+        dion_selection_scope=train_args.dion_selection_scope,
         scheduler_type="none",
         hidden_states_dtype=hidden_states_dtype,
         log_freq=1,

scripts/train.py

@@ -692,6 +692,8 @@ def main(cfg: TrainConfig):  # noqa: C901
         muon_weight_decay=args.muon_weight_decay,
         muon_ns_steps=args.muon_ns_steps,
         muon_adjust_lr_fn=args.muon_adjust_lr_fn,
+        dion_fraction=args.dion_fraction,
+        dion_selection_scope=args.dion_selection_scope,
         scheduler_type=args.scheduler_type,
         scheduler_warmup_steps=args.scheduler_warmup_steps,
         scheduler_warmup_ratio=args.scheduler_warmup_ratio,

src/speculators/train/config/schema.py

@@ -323,10 +323,12 @@ def _loss_parseable(cls, v: str | None) -> str | None:
 
 
 class OptimizerArgs(_Group):
-    optimizer: Literal["adamw", "muon"] = Field(
+    optimizer: Literal["adamw", "muon", "dion3"] = Field(
         default="muon",
         description="Optimizer to use. 'muon' applies Muon to 2D weight matrices and "
-        "AdamW to the remaining params (norms, biases, embeddings, lm_head).",
+        "AdamW to the remaining params (norms, biases, embeddings, lm_head). 'dion3' "
+        "substitutes Dion3 (microsoft/dion) for Muon on the same split; it requires "
+        "the optional `dion` package.",
     )
     lr: float = Field(default=1e-4, description="Learning rate (AdamW / base group).")
     weight_decay: float = Field(
@@ -352,6 +354,21 @@ class OptimizerArgs(_Group):
         default="match_rms_adamw",
         description="Muon LR adjustment. 'match_rms_adamw' matches AdamW's update RMS.",
     )
+    dion_fraction: float = Field(
+        default=0.25,
+        gt=0.0,
+        le=1.0,
+        description="Fraction of momentum-matrix rows Dion3 orthogonalizes per step. "
+        "This is the knob that buys Dion3 its speed; 1.0 disables compression and is "
+        "slower than Muon. Only used with --optimizer dion3.",
+    )
+    dion_selection_scope: Literal["local", "global"] = Field(
+        default="global",
+        description="How Dion3 picks rows when parameters are sharded. 'global' takes "
+        "the exact top-k on the assembled matrix, so the update does not depend on the "
+        "world size; 'local' takes a cheaper per-rank top-k whose result varies with "
+        "sharding. Only used with --optimizer dion3.",
+    )
 
 
 class SchedulerArgs(_Group):

src/speculators/train/optimizers.py

@@ -10,13 +10,23 @@
 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
 Newton-Schulz orthogonalization dispatches across ranks automatically.
+
+The "dion3" option substitutes Dion3 (``microsoft/dion``) for Muon over the identical
+parameter split. Dion3 orthogonalizes only a ``dion_fraction`` of the momentum matrix's
+rows per step and megabatches the sharded transfer, which is why its advantage grows
+with world size: Muon has to reassemble whole matrices from their shards to run
+Newton-Schulz, so sharding buys it nothing. ``dion`` is not published on PyPI and is
+therefore imported lazily rather than declared as a dependency; install it with
+``pip install git+https://github.com/microsoft/dion.git``.
 """
 
+import functools
 import logging
 
 import torch
 from torch import Tensor
 from torch.nn import Module
+from torch.torch_version import TorchVersion
 
 logger = logging.getLogger("speculators")
 
@@ -107,4 +117,124 @@ def build_optimizers(model: Module, config) -> list[torch.optim.Optimizer]:
             raise ValueError("No trainable parameters found to optimize.")
         return optimizers
 
+    if config.optimizer == "dion3":
+        return _build_dion3(model, config)
+
     raise ValueError(f"Unsupported optimizer: {config.optimizer!r}")
+
+
+def _device_mesh_of(params: list[Tensor]) -> object | None:
+    """Return the device mesh the parameters are sharded over, if any.
+
+    Taken from the parameters themselves rather than rebuilt with
+    ``init_device_mesh``, so it is by construction the same mesh ``fully_shard``
+    used. Without it Dion3 silently runs its single-device orthonormalization
+    path, which is where its multi-GPU advantage comes from.
+    """
+    for param in params:
+        mesh = getattr(param, "device_mesh", None)
+        if mesh is not None:
+            return mesh
+    return None
+
+
+def _static_shape_step(optimizer: torch.optim.Optimizer) -> torch.optim.Optimizer:
+    """Run ``optimizer.step`` with dynamic shapes disabled.
+
+    torch 2.13 regressed inductor's handling of Dion3's
+    ``@torch.compile(fullgraph=True)`` per-neuron normalization: once dynamo
+    promotes its shapes to dynamic -- which happens on the second distinct matrix
+    shape in a step -- the generated Triton kernel reuses a value emitted inside
+    the reduction loop from the epilogue after that loop closes. A ``tl.range``
+    body is a separate scope, so it fails to compile with
+    ``NameError: tmp<N> is not defined``. torch 2.12.1 and 2.12.0 emit a
+    self-contained epilogue and are unaffected, so this is gated to >= 2.13.
+
+    An optimizer's shapes are fixed by the parameter list, so there is nothing to
+    gain from dynamic shapes here. Scoping the setting to ``step`` rather than
+    setting it process-wide leaves the model's own ``torch.compile`` alone --
+    setting it globally measurably inflates the forward pass. Pinning shapes
+    static does cost extra recompiles, hence the raised ``cache_size_limit``
+    (cf. microsoft/dion#23).
+    """
+    if TorchVersion(torch.__version__) < (2, 13):
+        return optimizer
+
+    inner = optimizer.step
+
+    @functools.wraps(inner)
+    def step(*args, **kwargs):
+        with torch._dynamo.config.patch(  # noqa: SLF001
+            automatic_dynamic_shapes=False, cache_size_limit=64
+        ):
+            return inner(*args, **kwargs)
+
+    optimizer.step = step  # type: ignore[method-assign]
+    return optimizer
+
+
+def _build_dion3(model: Module, config) -> list[torch.optim.Optimizer]:
+    """Build a single Dion3 optimizer covering both parameter groups."""
+    try:
+        # Imported lazily on purpose: dion is not on PyPI, so it cannot be a
+        # declared dependency and may legitimately be absent.
+        from dion import Dion3  # noqa: PLC0415
+    except ImportError as exc:  # pragma: no cover - depends on optional install
+        raise ImportError(
+            "--optimizer dion3 requires the 'dion' package, which is not published "
+            "on PyPI. Install it with:\n"
+            "    pip install git+https://github.com/microsoft/dion.git"
+        ) from exc
+
+    matrix_named, scalar_named = split_named_params_for_muon(model)
+    if not matrix_named:
+        raise ValueError("No trainable 2D parameters found to optimize.")
+    matrix = [p for _, p in matrix_named]
+    scalar = [p for _, p in scalar_named]
+
+    logger.info(
+        "Dion3 optimizer: %d 2D params via Dion3 (fraction=%s, selection_scope=%s), "
+        "%d params via AdamW.",
+        len(matrix),
+        config.dion_fraction,
+        config.dion_selection_scope,
+        len(scalar),
+    )
+
+    # Dion3 consumes both groups itself, so this returns one optimizer where the
+    # muon path returns two. The trainer and checkpointer both take a list.
+    param_groups: list[dict] = [
+        {
+            "params": matrix,
+            "algorithm": "nordion2",
+            "lr": config.muon_lr,
+            "weight_decay": config.muon_weight_decay,
+        }
+    ]
+    if scalar:
+        param_groups.append(
+            {
+                "params": scalar,
+                "algorithm": "adamw",
+                "lr": config.lr,
+                "weight_decay": config.weight_decay,
+            }
+        )
+
+    optimizer = Dion3(
+        param_groups,
+        lr=config.muon_lr,
+        mu=config.muon_momentum,
+        # Preserve the torch.optim.AdamW defaults used by Muon's scalar group.
+        # Dion3 otherwise defaults beta2 to 0.95 instead of torch's 0.999.
+        betas=(0.9, 0.999),
+        weight_decay=config.muon_weight_decay,
+        fraction=config.dion_fraction,
+        # dion's "rms_norm" is 0.2*sqrt(max(fan_out, fan_in)), the same expression
+        # as torch Muon's "match_rms_adamw", so a given --muon-lr means the same
+        # effective step size on both. dion's own default is a different scale.
+        adjust_lr="rms_norm",
+        selection_scope=config.dion_selection_scope,
+        distributed_mesh=_device_mesh_of(matrix),
+    )
+    return [_static_shape_step(optimizer)]

src/speculators/train/trainer.py

@@ -105,13 +105,15 @@ class TrainerConfig(NamedTuple):
     resume_from_checkpoint: bool = False
     train_call_kwargs: dict | None = None
     val_call_kwargs: dict | None = None
-    optimizer: Literal["adamw", "muon"] = "adamw"
+    optimizer: Literal["adamw", "muon", "dion3"] = "adamw"
     weight_decay: float = 0.01
     muon_lr: float = 0.02
     muon_momentum: float = 0.95
     muon_weight_decay: float = 0.1
     muon_ns_steps: int = 5
     muon_adjust_lr_fn: str = "match_rms_adamw"
+    dion_fraction: float = 0.25
+    dion_selection_scope: str = "global"
     scheduler_type: Literal["linear", "cosine", "none"] = "linear"
     scheduler_warmup_steps: int | None = None
     scheduler_warmup_ratio: float | None = None

tests/unit/train/test_dion3_optimizer.py

@@ -0,0 +1,143 @@
+"""Dion3 is an opt-in substitute for Muon over the identical parameter split.
+
+These tests pin the parts that are easy to get wrong and cheap to check without a
+GPU or the optional ``dion`` install: that the option is accepted by the config,
+that the parameter split is the same one Muon uses, that the learning-rate
+convention is matched rather than left at dion's default, and that a missing
+``dion`` install fails with an actionable message rather than an ImportError from
+somewhere deep inside the trainer.
+"""
+
+import sys
+from types import SimpleNamespace
+from unittest import mock
+
+import pytest
+import torch
+from torch import nn
+
+from speculators.train.optimizers import (
+    build_optimizers,
+    split_named_params_for_muon,
+)
+
+
+class _Tiny(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.proj = nn.Linear(8, 16, bias=False)
+        self.norm = nn.LayerNorm(8)
+        self.embed_tokens = nn.Embedding(32, 8)
+
+
+def _config(**over):
+    base = {
+        "optimizer": "dion3",
+        "lr": 1e-4,
+        "weight_decay": 0.01,
+        "muon_lr": 1e-3,
+        "muon_momentum": 0.95,
+        "muon_weight_decay": 0.1,
+        "muon_ns_steps": 5,
+        "muon_adjust_lr_fn": "match_rms_adamw",
+        "dion_fraction": 0.25,
+        "dion_selection_scope": "global",
+    }
+    base.update(over)
+    return SimpleNamespace(**base)
+
+
+def test_config_accepts_dion3():
+    from speculators.train.config.schema import OptimizerArgs  # noqa: PLC0415
+
+    args = OptimizerArgs(optimizer="dion3", dion_fraction=0.5)
+    assert args.optimizer == "dion3"
+    assert args.dion_fraction == 0.5
+    # fraction is a proportion; 0 and >1 are meaningless
+    with pytest.raises(ValueError):
+        OptimizerArgs(optimizer="dion3", dion_fraction=0.0)
+    with pytest.raises(ValueError):
+        OptimizerArgs(optimizer="dion3", dion_fraction=1.5)
+
+
+def test_missing_dion_is_an_actionable_error():
+    """dion is not on PyPI, so this is the common failure and must explain itself."""
+    with (
+        mock.patch.dict(sys.modules, {"dion": None}),
+        pytest.raises(ImportError, match="github.com/microsoft/dion"),
+    ):
+        build_optimizers(_Tiny(), _config())
+
+
+def test_dion3_reuses_the_muon_parameter_split():
+    """The split must be Muon's, so the two optimizers stay comparable."""
+    model = _Tiny()
+    matrix, scalar = split_named_params_for_muon(model)
+    matrix_names = {n for n, _ in matrix}
+    scalar_names = {n for n, _ in scalar}
+
+    assert "proj.weight" in matrix_names
+    # embeddings and norms are excluded from the orthogonalized group
+    assert "embed_tokens.weight" in scalar_names
+    assert "norm.weight" in scalar_names
+    assert not matrix_names & scalar_names
+
+
+def test_dion3_is_constructed_with_matched_optimizer_conventions():
+    """Dion's defaults differ from the existing Muon and scalar AdamW paths.
+
+    ``rms_norm`` is ``0.2*sqrt(max(fan_out, fan_in))``, the same expression as
+    torch Muon's ``match_rms_adamw``. Getting this wrong silently changes the
+    effective learning rate, which would make any Muon-vs-Dion3 comparison
+    meaningless. Dion also defaults AdamW beta2 to 0.95, while the existing
+    scalar optimizer uses torch's 0.999 default.
+    """
+    pytest.importorskip("dion", reason="optional dependency, install from git")
+    captured = {}
+
+    import dion  # noqa: PLC0415
+
+    class _Spy(dion.Dion3):
+        def __init__(self, param_groups, **kwargs):
+            captured.update(kwargs)
+            captured["groups"] = [g.get("algorithm") for g in param_groups]
+            super().__init__(param_groups, **kwargs)
+
+    with mock.patch.object(dion, "Dion3", _Spy):
+        optimizers = build_optimizers(_Tiny(), _config())
+
+    assert len(optimizers) == 1
+    assert captured["adjust_lr"] == "rms_norm"
+    assert captured["betas"] == (0.9, 0.999)
+    assert captured["fraction"] == 0.25
+    assert captured["selection_scope"] == "global"
+    assert captured["groups"] == ["nordion2", "adamw"]
+    adamw_group = optimizers[0].param_groups[1]
+    assert (adamw_group["beta1"], adamw_group["beta2"]) == (0.9, 0.999)
+
+
+def test_step_is_wrapped_for_static_shapes():
+    """The workaround must be scoped, not process-global, and version-gated.
+
+    torch 2.13 regressed inductor here; 2.12.x is unaffected and should not pay
+    the extra recompiles that pinning shapes static costs.
+    """
+    pytest.importorskip("dion", reason="optional dependency, install from git")
+    dynamo_config = torch._dynamo.config
+    before = dynamo_config.automatic_dynamic_shapes
+    optimizers = build_optimizers(_Tiny(), _config())
+    # constructing must not touch the global flag
+    assert dynamo_config.automatic_dynamic_shapes is before
+    # and step must be the instance-level wrapper on the affected torch versions
+    from torch.torch_version import TorchVersion  # noqa: PLC0415
+
+    affected = TorchVersion(torch.__version__) >= (2, 13)
+    assert ("step" in optimizers[0].__dict__) is affected
+
+
+def test_muon_path_is_unchanged():
+    """The default path must not be perturbed by the new branch."""
+    optimizers = build_optimizers(_Tiny(), _config(optimizer="muon"))
+    assert [type(o).__name__ for o in optimizers] == ["Muon", "AdamW"]
+    # the dion3 step wrapper must not be applied to the muon path
+    assert "step" not in optimizers[0].__dict__