[feat] train: add Muon optimizer (FSDP2/DTensor-aware) alongside AdamW

fastvideo/tests/train/utils/test_muon.py

@@ -0,0 +1,98 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Unit tests for the Muon optimizer (CPU; no distributed required).
+
+Covers the three correctness surfaces of ``fastvideo/train/utils/muon.py``:
+the Newton-Schulz orthogonalization, the muon-vs-aux parameter split, and a
+single-optimizer step that reduces a quadratic loss with finite updates. The
+FSDP2 / DTensor gather-per-matrix path is exercised separately on multi-GPU
+hardware (not in CPU CI).
+"""
+
+from __future__ import annotations
+
+import torch
+
+from fastvideo.train.utils.muon import (
+    MuonWithAuxAdam,
+    split_params_for_muon,
+    zeropower_via_newtonschulz5,
+)
+
+
+def test_newton_schulz_orthogonalizes() -> None:
+    # The 5-step quintic pushes singular values into a band around 1 (it is
+    # deliberately approximate, not exact orthogonalization).
+    torch.manual_seed(0)
+    for shape in [(64, 128), (128, 64), (96, 96)]:
+        g = torch.randn(*shape)
+        o = zeropower_via_newtonschulz5(g, steps=5).float()
+        assert o.shape == g.shape
+        s = torch.linalg.svdvals(o)
+        assert torch.isfinite(s).all()
+        # The 5-step quintic lands singular values in a band around 1 (wider
+        # for square inputs); it is approximate by design, not exact.
+        assert 0.3 < s.min() and s.max() < 1.5
+
+
+def _named(specs: list[tuple[str, tuple[int, ...]]]):
+    return [(n, torch.nn.Parameter(torch.randn(*shp))) for n, shp in specs]
+
+
+def test_split_keeps_hidden_matmuls_excludes_embed_head_and_1d() -> None:
+    named = _named([
+        ("patch_embed.weight", (16, 3, 2, 2, 2)),       # conv (>2D) -> aux
+        ("blocks.0.attn.to_q.weight", (16, 16)),         # hidden 2D -> muon
+        ("blocks.0.attn.to_q.bias", (16, )),             # 1D -> aux
+        ("blocks.0.attn.to_out.weight", (16, 16)),       # attn out -> muon
+        ("blocks.0.mlp.fc1.weight", (64, 16)),           # hidden 2D -> muon
+        ("blocks.0.norm.weight", (16, )),                # 1D -> aux
+        ("text_embedder.weight", (16, 32)),              # embed -> aux
+        ("proj_out.weight", (3, 16)),                    # output head -> aux
+    ])
+    muon, aux = split_params_for_muon(named)
+    muon_names = {n for n, p in named if any(p is q for q in muon)}
+    aux_names = {n for n, p in named if any(p is q for q in aux)}
+    assert muon_names == {
+        "blocks.0.attn.to_q.weight",
+        "blocks.0.attn.to_out.weight",
+        "blocks.0.mlp.fc1.weight",
+    }
+    for excluded in ("patch_embed.weight", "text_embedder.weight",
+                     "proj_out.weight", "blocks.0.norm.weight",
+                     "blocks.0.attn.to_q.bias"):
+        assert excluded in aux_names
+
+
+def test_split_skips_non_trainable() -> None:
+    p = torch.nn.Parameter(torch.randn(8, 8))
+    p.requires_grad_(False)
+    muon, aux = split_params_for_muon([("blocks.0.w.weight", p)])
+    assert muon == [] and aux == []
+
+
+def test_muon_with_aux_adam_step_reduces_loss_and_is_finite() -> None:
+    torch.manual_seed(0)
+    w = torch.nn.Parameter(torch.randn(32, 32))   # muon group
+    b = torch.nn.Parameter(torch.zeros(32))       # aux (adam) group
+    opt = MuonWithAuxAdam([w], [b], lr=0.05, momentum=0.95, ns_steps=5,
+                          aux_lr=0.02)
+    tw, tb = torch.randn(32, 32), torch.randn(32)
+    losses = []
+    for _ in range(100):
+        opt.zero_grad()
+        loss = ((w - tw)**2).mean() + ((b - tb)**2).mean()
+        loss.backward()
+        opt.step()
+        losses.append(loss.item())
+    assert torch.isfinite(w).all() and torch.isfinite(b).all()
+    # Both groups must make progress (Muon orthogonalizes the update, so it is
+    # steadier than GD but still monotonically descends this convex problem).
+    assert losses[-1] < 0.5 * losses[0]
+
+
+def test_muon_requires_some_params() -> None:
+    try:
+        MuonWithAuxAdam([], [], lr=0.01)
+    except ValueError:
+        return
+    raise AssertionError("MuonWithAuxAdam([], []) should raise ValueError")

fastvideo/train/methods/distribution_matching/dmd2.py

@@ -320,6 +320,7 @@ def _init_optimizers_and_schedulers(self) -> None:
             self._student_lr_scheduler,
         ) = build_optimizer_and_scheduler(
             params=student_params,
+            module=self.student.transformer,
             optimizer_config=tc.optimizer,
             loop_config=tc.loop,
             learning_rate=student_lr,
@@ -359,6 +360,7 @@ def _init_optimizers_and_schedulers(self) -> None:
             self._critic_lr_scheduler,
         ) = build_optimizer_and_scheduler(
             params=critic_params,
+            module=self.critic.transformer,
             optimizer_config=tc.optimizer,
             loop_config=tc.loop,
             learning_rate=critic_lr,

fastvideo/train/methods/fine_tuning/dfsft.py

@@ -305,6 +305,7 @@ def _init_optimizers_and_schedulers(self) -> None:
             self._student_lr_scheduler,
         ) = build_optimizer_and_scheduler(
             params=student_params,
+            module=self.student.transformer,
             optimizer_config=tc.optimizer,
             loop_config=tc.loop,
             learning_rate=student_lr,

fastvideo/train/methods/fine_tuning/finetune.py

@@ -169,6 +169,7 @@ def _init_optimizers_and_schedulers(self) -> None:
             self._student_lr_scheduler,
         ) = build_optimizer_and_scheduler(
             params=student_params,
+            module=self.student.transformer,
             optimizer_config=tc.optimizer,
             loop_config=tc.loop,
             learning_rate=student_lr,

fastvideo/train/methods/knowledge_distillation/kd.py

@@ -756,6 +756,7 @@ def _init_optimizers_and_schedulers(self) -> None:
             self._student_lr_scheduler,
         ) = build_optimizer_and_scheduler(
             params=student_params,
+            module=self.student.transformer,
             optimizer_config=tc.optimizer,
             loop_config=tc.loop,
             learning_rate=student_lr,

fastvideo/train/methods/rl/diffusion_nft.py

@@ -266,6 +266,7 @@ def _init_optimizer_and_scheduler(self) -> None:
         params = [p for p in self.student.transformer.parameters() if p.requires_grad]
         self._student_optimizer, self._student_lr_scheduler = build_optimizer_and_scheduler(
             params=params,
+            module=self.student.transformer,
             optimizer_config=self.training_config.optimizer,
             loop_config=self.training_config.loop,
             learning_rate=lr,

fastvideo/train/utils/config.py

@@ -379,6 +379,11 @@ def _build_training_config(
             lr_num_cycles=int(o.get("lr_num_cycles", 0) or 0),
             lr_power=float(o.get("lr_power", 0.0) or 0.0),
             min_lr_ratio=float(o.get("min_lr_ratio", 0.5) or 0.5),
+            optimizer_type=str(o.get("optimizer_type", "adamw") or "adamw").strip().lower(),
+            muon_lr=float(o.get("muon_lr", 0.0) or 0.0),
+            muon_momentum=float(o.get("muon_momentum", 0.95) or 0.95),
+            muon_weight_decay=float(o.get("muon_weight_decay", 0.0) or 0.0),
+            muon_ns_steps=int(o.get("muon_ns_steps", 5) or 5),
         ),
         loop=TrainingLoopConfig(
             max_train_steps=int(lo.get("max_train_steps", 0) or 0),

fastvideo/train/utils/muon.py

@@ -0,0 +1,209 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Muon optimizer (FSDP2 / DTensor-aware) with an auxiliary AdamW group.
+
+Muon (MomentUm Orthogonalized by Newton-schulz, Keller Jordan 2024) replaces
+the raw momentum update of a 2-D weight matrix with its nearest semi-orthogonal
+matrix, computed by a few Newton-Schulz iterations. It is applied to the hidden
+weight matrices of the transformer (attention q/k/v/out projections, MLP fc
+layers) while embeddings, the output head, and all 1-D parameters (norm/bias)
+stay on AdamW — the standard Muon recipe.
+
+FSDP2 note: under ``fully_shard`` the parameters are ``DTensor``s sharded along
+dim-0, but Newton-Schulz needs the *full* 2-D matrix. We keep the momentum
+buffer sharded (memory-efficient, like the parameter) and only gather each
+matrix transiently for the orthogonalization step (``full_tensor()``), then
+re-shard the update back to the parameter's placement (gather-per-matrix). Peak
+extra memory is one full matrix at a time, not the whole replicated optimizer
+state.
+"""
+
+from __future__ import annotations
+
+import torch
+
+try:
+    from torch.distributed.tensor import DTensor, distribute_tensor
+    _HAS_DTENSOR = True
+except Exception:  # pragma: no cover - older torch
+    DTensor = ()  # type: ignore[assignment]
+    _HAS_DTENSOR = False
+
+
+def zeropower_via_newtonschulz5(G: torch.Tensor, steps: int) -> torch.Tensor:
+    """Newton-Schulz iteration to orthogonalize ``G`` (Keller Jordan quintic).
+
+    Returns a matrix with the same shape as ``G`` whose singular values are
+    pushed toward 1. Runs in bf16; the quintic coefficients are tuned so the
+    iteration converges from a spectral-norm-normalized start.
+    """
+    assert G.ndim == 2, f"expected a 2-D matrix, got shape {tuple(G.shape)}"
+    a, b, c = (3.4445, -4.7750, 2.0315)
+    X = G.bfloat16()
+    transposed = False
+    if X.size(0) > X.size(1):
+        X = X.T
+        transposed = True
+    # Normalize so the spectral norm is <= 1 before the iteration.
+    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 transposed:
+        X = X.T
+    return X.to(G.dtype)
+
+
+def _is_dtensor(t: torch.Tensor) -> bool:
+    return _HAS_DTENSOR and isinstance(t, DTensor)
+
+
+def _full(t: torch.Tensor) -> torch.Tensor:
+    """Gather a (possibly sharded) tensor to its full local form."""
+    return t.full_tensor() if _is_dtensor(t) else t
+
+
+def _like_param(update_full: torch.Tensor, param: torch.Tensor) -> torch.Tensor:
+    """Re-shard a full update tensor to match ``param``'s DTensor placement."""
+    if _is_dtensor(param):
+        return distribute_tensor(update_full, param.device_mesh, param.placements)
+    return update_full
+
+
+def split_params_for_muon(
+    named_params: list[tuple[str, torch.nn.Parameter]], ) -> tuple[list[torch.nn.Parameter], list[torch.nn.Parameter]]:
+    """Partition trainable params into (muon, aux-adam) groups.
+
+    Muon eligibility: exactly 2-D weight matrices that are *not* embeddings or
+    the final output head. Everything else — 1-D params (norms/biases), >2-D
+    conv/patch-embed weights, embeddings, and the output projection — goes to
+    AdamW. ``out_proj`` / ``to_out`` (attention output projections) ARE hidden
+    matmuls and stay on Muon; only the model's final head (``proj_out`` /
+    ``final_layer`` / ``*head``) is excluded.
+    """
+    _EXCLUDE = ("embed", "embedder", "proj_out", "final_layer", "head")
+    muon: list[torch.nn.Parameter] = []
+    aux: list[torch.nn.Parameter] = []
+    for name, p in named_params:
+        if not p.requires_grad:
+            continue
+        lname = name.lower()
+        if p.ndim == 2 and not any(tok in lname for tok in _EXCLUDE):
+            muon.append(p)
+        else:
+            aux.append(p)
+    return muon, aux
+
+
+class MuonWithAuxAdam(torch.optim.Optimizer):
+    """Single optimizer running Muon on group 0 and AdamW on the aux group.
+
+    Param groups carry a ``use_muon`` flag. The Muon group uses
+    ``lr/momentum/weight_decay/ns_steps``; the aux group uses
+    ``lr/betas/eps/weight_decay``. Keeping both in one optimizer means the
+    existing single-optimizer / single-scheduler trainer plumbing is unchanged.
+    """
+
+    def __init__(
+        self,
+        muon_params: list[torch.nn.Parameter],
+        aux_params: list[torch.nn.Parameter],
+        *,
+        lr: float,
+        momentum: float = 0.95,
+        weight_decay: float = 0.0,
+        ns_steps: int = 5,
+        nesterov: bool = True,
+        aux_lr: float | None = None,
+        aux_betas: tuple[float, float] = (0.9, 0.999),
+        aux_eps: float = 1e-8,
+        aux_weight_decay: float | None = None,
+    ) -> None:
+        groups = []
+        if muon_params:
+            groups.append({
+                "params": muon_params,
+                "use_muon": True,
+                "lr": float(lr),
+                "momentum": float(momentum),
+                "weight_decay": float(weight_decay),
+                "ns_steps": int(ns_steps),
+                "nesterov": bool(nesterov),
+            })
+        if aux_params:
+            groups.append({
+                "params": aux_params,
+                "use_muon": False,
+                "lr": float(aux_lr if aux_lr is not None else lr),
+                "betas": tuple(aux_betas),
+                "eps": float(aux_eps),
+                "weight_decay": float(aux_weight_decay if aux_weight_decay is not None else weight_decay),
+            })
+        if not groups:
+            raise ValueError("MuonWithAuxAdam got no parameters")
+        super().__init__(groups, {})
+
+    @torch.no_grad()
+    def step(self, closure=None):  # type: ignore[override]
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+        for group in self.param_groups:
+            if group["use_muon"]:
+                self._muon_step(group)
+            else:
+                self._adam_step(group)
+        return loss
+
+    def _muon_step(self, group) -> None:
+        lr = group["lr"]
+        momentum = group["momentum"]
+        wd = group["weight_decay"]
+        ns_steps = group["ns_steps"]
+        nesterov = group["nesterov"]
+        for p in group["params"]:
+            if p.grad is None:
+                continue
+            g = p.grad
+            state = self.state[p]
+            if "momentum_buffer" not in state:
+                state["momentum_buffer"] = torch.zeros_like(p)
+            buf = state["momentum_buffer"]
+            # Sharded momentum update (DTensor ops stay on local shards).
+            buf.mul_(momentum).add_(g)
+            g_eff = g.add(buf, alpha=momentum) if nesterov else buf
+            # Gather the full matrix only for the orthogonalization.
+            update_full = zeropower_via_newtonschulz5(_full(g_eff), ns_steps)
+            # Keller-Jordan RMS-matching scale for non-square matrices.
+            rows, cols = update_full.shape
+            scale = max(1.0, rows / cols)**0.5
+            if wd != 0.0:
+                p.mul_(1.0 - lr * wd)
+            p.add_(_like_param(update_full, p), alpha=-lr * scale)
+
+    def _adam_step(self, group) -> None:
+        lr = group["lr"]
+        beta1, beta2 = group["betas"]
+        eps = group["eps"]
+        wd = group["weight_decay"]
+        for p in group["params"]:
+            if p.grad is None:
+                continue
+            g = p.grad
+            state = self.state[p]
+            if "step" not in state:
+                state["step"] = 0
+                state["exp_avg"] = torch.zeros_like(p)
+                state["exp_avg_sq"] = torch.zeros_like(p)
+            state["step"] += 1
+            t = state["step"]
+            exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
+            exp_avg.mul_(beta1).add_(g, alpha=1.0 - beta1)
+            exp_avg_sq.mul_(beta2).addcmul_(g, g, value=1.0 - beta2)
+            bias1 = 1.0 - beta1**t
+            bias2 = 1.0 - beta2**t
+            denom = (exp_avg_sq.sqrt() / (bias2**0.5)).add_(eps)
+            if wd != 0.0:
+                p.mul_(1.0 - lr * wd)
+            p.addcdiv_(exp_avg, denom, value=-lr / bias1)