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

[alexzms]

Purpose

fastvideo/train only offered AdamW. This adds Muon (MomentUm Orthogonalized by Newton-schulz, Keller Jordan 2024) as an opt-in optimizer for the modular trainer.

How it works

Muon replaces a 2-D weight matrix's raw momentum update with its nearest semi-orthogonal matrix (a few Newton-Schulz iterations). Following the standard recipe, it's applied to the transformer's 2-D hidden weight matrices (attention q/k/v/out projections, MLP fc layers) while embeddings, the output head, and all 1-D params (norm/bias) fall back to an auxiliary AdamW group. Both live in one MuonWithAuxAdam optimizer (two param groups), so the existing single-optimizer / single-scheduler trainer plumbing is unchanged.

FSDP2 / DTensor: under fully_shard the params are DTensors sharded on dim-0, but Newton-Schulz needs the full 2-D matrix. The momentum buffer stays sharded (memory-efficient); each matrix is gathered transiently (full_tensor()) only for the orthogonalization, then the update is re-sharded back to the param's placement (gather-per-matrix). Peak extra memory is one full matrix at a time, not the whole replicated optimizer state. (A fully-distributed Newton-Schulz that avoids the per-matrix gather is a natural follow-up.)

Usage

training:
  optimizer:
    optimizer_type: muon        # default "adamw"
    learning_rate: 1.0e-4       # aux-AdamW group (embeddings/head/1-D)
    muon_lr: 2.0e-2             # Muon group (0 -> reuse learning_rate)
    muon_momentum: 0.95
    muon_ns_steps: 5

Changes

• fastvideo/train/utils/muon.py (new): zeropower_via_newtonschulz5, split_params_for_muon, MuonWithAuxAdam.
• OptimizerConfig: optimizer_type + muon_lr/muon_momentum/muon_weight_decay/muon_ns_steps; parsed in config.py.
• build_optimizer_and_scheduler: branch on optimizer_type; the muon path takes module to split params by name. All 5 method callers pass module=.
• fastvideo/tests/train/utils/test_muon.py (new): unit tests.

Test evidence

• CPU unit tests (5 passed): NS orthogonalization (singular values band around 1), param split (attention out-proj stays Muon; embeddings/head/1-D -> aux), MuonWithAuxAdam reduces a convex loss with finite updates, empty-params guard.
• 2-GPU FSDP2 run: 20 steps through the DTensor gather-per-matrix path, all params DTensors (2 muon / 6 aux), loss 0.115 -> 0.033, finite — no error.

Test plan

• pytest fastvideo/tests/train/utils/test_muon.py (CPU) -> 5 passed
• 2-GPU FSDP2 smoke (DTensor gather-per-matrix)
• /test train-framework on CI

[mergify]

Merge Protections

🔴 1 of 1 protections blocking · waiting on 👀 reviews and 🤖 CI

	Protection	Waiting on
🔴	PR merge requirements	👀 reviews and 🤖 CI

🔴 PR merge requirements

Waiting for

• #approved-reviews-by>=1
• check-success=fastcheck-passed
• check-success=full-suite-passed
• check-success~=pre-commit

This rule is failing.

• #approved-reviews-by>=1
• check-success=fastcheck-passed
• check-success=full-suite-passed
• check-success~=pre-commit
• title~=(?i)^\[(feat|feature|bugfix|fix|refactor|perf|ci|doc|docs|misc|chore|kernel|new.?model|skill|skills|infra)\]

[gemini-code-assist]

Code Review

This pull request introduces the Muon optimizer with an auxiliary AdamW group (MuonWithAuxAdam) to optimize 2-D hidden weight matrices while keeping other parameters on AdamW, including configuration parsing, integration into various training methods, and unit tests. Feedback on the changes highlights several important improvements: preventing potential overflow in zeropower_via_newtonschulz5 by computing the norm in float32, ensuring only the parameters in the provided params list are optimized to support parameter-efficient fine-tuning, making the transposed tensor contiguous to avoid issues in distributed environments, and correctly handling negative values when defaulting muon_lr to learning_rate.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

In zeropower_via_newtonschulz5, calling X.norm() on a large bfloat16 tensor can easily overflow to inf because the intermediate sum of squares can exceed the maximum representable value in bfloat16 (65504). To prevent this, compute the norm in float32 before dividing.

Suggested change

	X = X / (X.norm() + 1e-7)
	X = X / (X.float().norm() + 1e-7)

[gemini-code-assist]

When optimizer_type == "muon", the optimizer parameters are extracted directly from module.named_parameters(). However, if the caller passed a specific subset of parameters in the params argument (e.g., for parameter-efficient fine-tuning like LoRA, or layer-wise learning rates), this will ignore that subset and optimize all trainable parameters in the module. Filter the parameters to ensure only those present in the params list are optimized.

        param_set = set(params)\n        muon_params, aux_params = split_params_for_muon([(n, p) for n, p in module.named_parameters()\n                                                         if p in param_set])

[gemini-code-assist]

When transposed is True, X.T returns a non-contiguous tensor. Some distributed operations (such as distribute_tensor or other collective communication methods) require contiguous memory layouts and will raise a runtime error if passed a non-contiguous tensor. Ensure the returned tensor is contiguous.

    if transposed:\n        X = X.T.contiguous()

[gemini-code-assist]

The docstring states that muon_lr defaults to learning_rate when <= 0. However, using float(optimizer_config.muon_lr) or float(learning_rate) will only default to learning_rate if muon_lr is exactly 0.0 (since negative values are truthy in Python). Explicitly check if muon_lr <= 0.0 to handle negative values correctly.

        muon_lr = float(optimizer_config.muon_lr)\n        if muon_lr <= 0.0:\n            muon_lr = float(learning_rate)