Make Nora+NAdamW the default optimizer

README.md

@@ -101,7 +101,9 @@ nohup python scripts/tools/train.py \
     --batch_per_device=4 \
     --grad_accum_steps=8 \
     --warmup_ratio=0.05 \
+    --optimizer=nora_nadamw \
     --max_lr=1e-4 \
+    --nora_lr=5e-3 \
     --max_grad_norm=1.0 \
     --weight_decay=0.1 \
     --window_size=2048 \

astrai/config/train_config.py

@@ -30,6 +30,8 @@ class TrainConfig(BaseConfig):
         strategy (str): Training strategy (seq, sft, dpo, grpo, online_*).
         dataset (Dataset): Dataset for training.
         optimizer_fn (Callable[[nn.Module], Optimizer]): Optimizer factory for training.
+        optimizer_name (Optional[str]): Serializable built-in optimizer identifier. Defaults to None.
+        optimizer_hyperparameters (Dict[str, Any]): Serializable optimizer settings. Defaults to {}.
         scheduler_fn (Callable[[Optimizer], LRScheduler]): Scheduler factory for training.
         n_epoch (int): Number of epochs for training. Defaults to 1.
         batch_per_device (int): Batch size per device. Defaults to 4.
@@ -74,6 +76,8 @@ class TrainConfig(BaseConfig):
     dataset: Dataset
     optimizer_fn: Callable[[nn.Module], Optimizer]
     scheduler_fn: Callable[[Optimizer], LRScheduler]
+    optimizer_name: Optional[str] = None
+    optimizer_hyperparameters: Dict[str, Any] = field(default_factory=dict)
     n_epoch: int = 1
     batch_per_device: int = 4
     grad_accum_steps: int = 1

astrai/optim/__init__.py

@@ -0,0 +1,36 @@
+"""Optimizer implementations and factory registration."""
+
+from torch.optim import Optimizer
+
+from astrai.factory import BaseFactory
+
+
+class OptimizerFactory(BaseFactory[Optimizer]):
+    """Factory for built-in training optimizers."""
+
+
+from astrai.optim.muon_mix import MuonMix
+from astrai.optim.nora_nadamw import (
+    NAdamW,
+    Nora,
+    NoraNAdamW,
+    OptimizerParameterGroups,
+    nora_direction,
+    nora_lr_scale,
+    partition_optimizer_parameters,
+)
+
+OptimizerFactory.register("nora_nadamw")(NoraNAdamW)
+OptimizerFactory.register("muon_adamw")(MuonMix)
+
+__all__ = [
+    "MuonMix",
+    "NAdamW",
+    "Nora",
+    "NoraNAdamW",
+    "OptimizerFactory",
+    "OptimizerParameterGroups",
+    "nora_direction",
+    "nora_lr_scale",
+    "partition_optimizer_parameters",
+]

astrai/optim/muon_mix.py

@@ -0,0 +1,91 @@
+"""Legacy Muon + AdamW combined optimizer."""
+
+from typing import Any
+
+import torch
+from torch import Tensor, nn, optim
+
+
+class MuonMix(optim.Optimizer):
+    """Combined Muon (matrix) + AdamW (non-matrix) optimizer."""
+
+    optimizer_name = "muon_adamw"
+
+    def __init__(
+        self,
+        model: nn.Module,
+        lr: float = 3e-4,
+        weight_decay: float = 0.1,
+        momentum: float = 0.95,
+        nesterov: bool = True,
+        ns_steps: int = 5,
+        adjust_lr_fn: str = "match_rms_adamw",
+    ):
+        defaults = {
+            "lr": lr,
+            "weight_decay": weight_decay,
+            "momentum": momentum,
+            "nesterov": nesterov,
+            "ns_steps": ns_steps,
+            "adjust_lr_fn": adjust_lr_fn,
+        }
+        params = [param for param in model.parameters() if param.requires_grad]
+        super().__init__(params, defaults)
+
+        matrix_params: list[Tensor] = []
+        other_params: list[Tensor] = []
+        for name, param in model.named_parameters():
+            if not param.requires_grad:
+                continue
+            if (
+                param.dim() >= 2
+                and "norm" not in name
+                and "bias" not in name
+                and "embed" not in name
+                and "lm_head" not in name
+            ):
+                matrix_params.append(param)
+            else:
+                other_params.append(param)
+
+        self.muon = optim.Muon(
+            matrix_params,
+            lr=lr,
+            weight_decay=weight_decay,
+            momentum=momentum,
+            nesterov=nesterov,
+            ns_steps=ns_steps,
+            adjust_lr_fn=adjust_lr_fn,
+        )
+        self.adamw = optim.AdamW(
+            [{"params": other_params, "weight_decay": 0.0}],
+            lr=lr,
+            betas=(0.9, 0.95),
+            fused=True,
+        )
+
+        self.param_groups = [*self.muon.param_groups, *self.adamw.param_groups]
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        self.muon.step(closure)
+        self.adamw.step(closure)
+
+    def zero_grad(self, set_to_none: bool = True):
+        self.muon.zero_grad(set_to_none=set_to_none)
+        self.adamw.zero_grad(set_to_none=set_to_none)
+
+    def state_dict(self) -> dict[str, Any]:
+        return {
+            "muon": self.muon.state_dict(),
+            "adamw": self.adamw.state_dict(),
+        }
+
+    def load_state_dict(self, state_dict: dict[str, Any]):
+        if "muon" not in state_dict or "adamw" not in state_dict:
+            raise ValueError(
+                "Checkpoint optimizer state is not compatible with muon_adamw"
+            )
+        self.muon.load_state_dict(state_dict["muon"])
+        self.adamw.load_state_dict(state_dict["adamw"])
+        self.param_groups = [*self.muon.param_groups, *self.adamw.param_groups]

astrai/optim/nora_nadamw.py

@@ -0,0 +1,379 @@
+"""Nora matrix optimizer combined with Nesterov AdamW."""
+
+import math
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from torch import Tensor, nn
+from torch.distributed.tensor import DTensor, Shard
+from torch.optim import Optimizer
+
+from astrai.model.components.embedding import Embedding
+from astrai.model.components.linear import Linear
+from astrai.model.components.lora import LoRALinear
+from astrai.model.components.norm import RMSNorm
+
+NORA_EPS = 1e-10
+
+
+def _row_normalize(tensor: Tensor, eps: float) -> Tensor:
+    return tensor / tensor.norm(dim=-1, keepdim=True).clamp(min=eps)
+
+
+def nora_direction(update: Tensor, param: Tensor, eps: float = NORA_EPS) -> Tensor:
+    """Project an update onto each parameter row's tangent space and normalize."""
+    theta_hat = _row_normalize(param.to(torch.float32), eps)
+    update_fp32 = update.to(torch.float32)
+    radial = (update_fp32 * theta_hat).sum(dim=-1, keepdim=True) * theta_hat
+    direction = _row_normalize(update_fp32 - radial, eps)
+    return direction.to(update.dtype)
+
+
+def nora_lr_scale(lr: float, shape: torch.Size) -> float:
+    """Scale Nora's LR for tall ``[d_out, d_in]`` linear weights."""
+    return lr * math.sqrt(max(1.0, shape[-2] / shape[-1]))
+
+
+def _validate_complete_rows(param: Tensor) -> None:
+    if not isinstance(param, DTensor):
+        return
+    last_dim = param.ndim - 1
+    for placement in param.placements:
+        if isinstance(placement, Shard) and placement.dim % param.ndim == last_dim:
+            raise ValueError(
+                "Nora requires complete parameter rows, but this DTensor is sharded "
+                "along its last dimension"
+            )
+
+
+class Nora(Optimizer):
+    """Normalized Orthogonal Row Alignment for two-dimensional matrices."""
+
+    def __init__(
+        self,
+        params,
+        lr: float = 5e-3,
+        weight_decay: float = 0.0,
+        momentum: float = 0.95,
+        beta: float = 0.95,
+        nesterov: bool = True,
+        eps: float = NORA_EPS,
+    ):
+        if lr < 0:
+            raise ValueError(f"Invalid learning rate: {lr}")
+        if weight_decay < 0:
+            raise ValueError(f"Invalid weight decay: {weight_decay}")
+        if not 0 <= momentum <= 1:
+            raise ValueError(f"Invalid momentum: {momentum}")
+        if not 0 <= beta < 1:
+            raise ValueError(f"Invalid beta: {beta}")
+        if eps <= 0:
+            raise ValueError(f"Invalid epsilon: {eps}")
+
+        defaults = {
+            "lr": lr,
+            "weight_decay": weight_decay,
+            "momentum": momentum,
+            "beta": beta,
+            "nesterov": nesterov,
+            "eps": eps,
+        }
+        super().__init__(params, defaults)
+        for group in self.param_groups:
+            for param in group["params"]:
+                if param.ndim != 2:
+                    raise ValueError(
+                        f"Nora only supports 2D matrices, got shape {tuple(param.shape)}"
+                    )
+                _validate_complete_rows(param)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            lr = group["lr"]
+            weight_decay = group["weight_decay"]
+            momentum = group["momentum"]
+            beta = group["beta"]
+            nesterov = group["nesterov"]
+            eps = group["eps"]
+            for param in group["params"]:
+                if param.grad is None:
+                    continue
+                if param.grad.is_sparse:
+                    raise RuntimeError("Nora does not support sparse gradients")
+
+                grad = param.grad
+                state = self.state[param]
+                momentum_buffer = state.get("momentum_buffer")
+                if momentum_buffer is None:
+                    momentum_buffer = torch.zeros_like(grad)
+                momentum_buffer.lerp_(grad, 1 - beta)
+                update = (
+                    grad.lerp(momentum_buffer, momentum)
+                    if nesterov
+                    else momentum_buffer
+                )
+                direction = nora_direction(update, param, eps)
+
+                if weight_decay != 0:
+                    param.mul_(1 - lr * weight_decay)
+                param.add_(direction, alpha=-nora_lr_scale(lr, param.shape))
+                state["momentum_buffer"] = momentum_buffer
+
+        return loss
+
+
+class NAdamW(Optimizer):
+    """AdamW using the reference Nesterov first-moment update."""
+
+    def __init__(
+        self,
+        params,
+        lr: float = 3e-4,
+        betas: tuple[float, float] = (0.9, 0.999),
+        eps: float = 1e-8,
+        weight_decay: float = 0.1,
+    ):
+        beta1, beta2 = betas
+        if lr < 0:
+            raise ValueError(f"Invalid learning rate: {lr}")
+        if not 0 <= beta1 < 1 or not 0 <= beta2 < 1:
+            raise ValueError(f"Invalid betas: {betas}")
+        if eps <= 0:
+            raise ValueError(f"Invalid epsilon: {eps}")
+        if weight_decay < 0:
+            raise ValueError(f"Invalid weight decay: {weight_decay}")
+        defaults = {
+            "lr": lr,
+            "betas": betas,
+            "eps": eps,
+            "weight_decay": weight_decay,
+        }
+        super().__init__(params, defaults)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            beta1, beta2 = group["betas"]
+            eps = group["eps"]
+            lr = group["lr"]
+            weight_decay = group["weight_decay"]
+            for param in group["params"]:
+                if param.grad is None:
+                    continue
+                if param.grad.is_sparse:
+                    raise RuntimeError("NAdamW does not support sparse gradients")
+
+                grad = param.grad
+                state = self.state[param]
+                if not state:
+                    state["step"] = 0
+                    state["m"] = torch.zeros_like(param)
+                    state["v"] = torch.zeros_like(param)
+
+                state["step"] += 1
+                first_moment = state["m"]
+                second_moment = state["v"]
+                first_moment.mul_(beta1).add_(grad, alpha=1 - beta1)
+                second_moment.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
+
+                bias_correction1 = 1 - beta1 ** state["step"]
+                bias_correction2 = 1 - beta2 ** state["step"]
+                nesterov_moment = (
+                    beta1 * first_moment + (1 - beta1) * grad
+                ) / bias_correction1
+                corrected_second_moment = second_moment / bias_correction2
+
+                if weight_decay != 0:
+                    param.mul_(1 - lr * weight_decay)
+                param.addcdiv_(
+                    nesterov_moment,
+                    corrected_second_moment.sqrt().add_(eps),
+                    value=-lr,
+                )
+
+        return loss
+
+
+@dataclass
+class OptimizerParameterGroups:
+    nora: list[Tensor]
+    nadamw_decay: list[Tensor]
+    nadamw_no_decay: list[Tensor]
+
+
+def partition_optimizer_parameters(model: nn.Module) -> OptimizerParameterGroups:
+    """Partition trainable parameters by module role and parameter identity."""
+    nora_ids: set[int] = set()
+    no_decay_ids: set[int] = set()
+
+    for module_name, module in model.named_modules():
+        if isinstance(module, LoRALinear):
+            for param in module.parameters(recurse=False):
+                if param.requires_grad:
+                    no_decay_ids.add(id(param))
+            continue
+
+        if isinstance(module, (Embedding, RMSNorm)):
+            for param in module.parameters(recurse=False):
+                if param.requires_grad:
+                    no_decay_ids.add(id(param))
+            continue
+
+        if not isinstance(module, Linear):
+            continue
+
+        if module.bias is not None and module.bias.requires_grad:
+            no_decay_ids.add(id(module.bias))
+        if not module.weight.requires_grad:
+            continue
+        if module_name.rsplit(".", 1)[-1] == "lm_head":
+            no_decay_ids.add(id(module.weight))
+        elif module.weight.ndim == 2:
+            nora_ids.add(id(module.weight))
+
+    nora: list[Tensor] = []
+    nadamw_decay: list[Tensor] = []
+    nadamw_no_decay: list[Tensor] = []
+    seen: set[int] = set()
+    for param in model.parameters():
+        param_id = id(param)
+        if not param.requires_grad or param_id in seen:
+            continue
+        seen.add(param_id)
+        if param_id in no_decay_ids or param.ndim <= 1:
+            nadamw_no_decay.append(param)
+        elif param_id in nora_ids:
+            nora.append(param)
+        else:
+            nadamw_decay.append(param)
+
+    trainable_ids = {id(param) for param in model.parameters() if param.requires_grad}
+    grouped_ids = {id(param) for param in [*nora, *nadamw_decay, *nadamw_no_decay]}
+    if grouped_ids != trainable_ids:
+        missing = len(trainable_ids - grouped_ids)
+        extra = len(grouped_ids - trainable_ids)
+        raise RuntimeError(
+            f"Optimizer parameter partition is incomplete: missing={missing}, extra={extra}"
+        )
+
+    return OptimizerParameterGroups(nora, nadamw_decay, nadamw_no_decay)
+
+
+class NoraNAdamW(Optimizer):
+    """Nora for internal linear weights and NAdamW for remaining parameters."""
+
+    optimizer_name = "nora_nadamw"
+
+    def __init__(
+        self,
+        model: nn.Module,
+        lr: float = 3e-4,
+        weight_decay: float = 0.1,
+        nora_lr: float = 5e-3,
+        nora_weight_decay: float = 0.0,
+        nora_beta: float = 0.95,
+        nora_momentum: float = 0.95,
+    ):
+        groups = partition_optimizer_parameters(model)
+        all_params = [
+            *groups.nora,
+            *groups.nadamw_decay,
+            *groups.nadamw_no_decay,
+        ]
+        if not all_params:
+            raise ValueError(
+                "Cannot build an optimizer for a model with no trainable parameters"
+            )
+        super().__init__(all_params, {})
+
+        self.nora = (
+            Nora(
+                groups.nora,
+                lr=nora_lr,
+                weight_decay=nora_weight_decay,
+                momentum=nora_momentum,
+                beta=nora_beta,
+            )
+            if groups.nora
+            else None
+        )
+
+        nadamw_groups = []
+        if groups.nadamw_decay:
+            nadamw_groups.append(
+                {"params": groups.nadamw_decay, "weight_decay": weight_decay}
+            )
+        if groups.nadamw_no_decay:
+            nadamw_groups.append(
+                {"params": groups.nadamw_no_decay, "weight_decay": 0.0}
+            )
+        self.nadamw = NAdamW(nadamw_groups, lr=lr) if nadamw_groups else None
+        self._refresh_param_groups()
+
+    def _refresh_param_groups(self) -> None:
+        self.param_groups = []
+        if self.nora is not None:
+            self.param_groups.extend(self.nora.param_groups)
+        if self.nadamw is not None:
+            self.param_groups.extend(self.nadamw.param_groups)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+        if self.nora is not None:
+            self.nora.step()
+        if self.nadamw is not None:
+            self.nadamw.step()
+        return loss
+
+    def zero_grad(self, set_to_none: bool = True):
+        if self.nora is not None:
+            self.nora.zero_grad(set_to_none=set_to_none)
+        if self.nadamw is not None:
+            self.nadamw.zero_grad(set_to_none=set_to_none)
+
+    def state_dict(self) -> dict[str, Any]:
+        return {
+            "nora": self.nora.state_dict() if self.nora is not None else None,
+            "nadamw": self.nadamw.state_dict() if self.nadamw is not None else None,
+        }
+
+    def load_state_dict(self, state_dict: dict[str, Any]):
+        if "muon" in state_dict or "adamw" in state_dict:
+            raise ValueError(
+                "Checkpoint uses muon_adamw state; select optimizer='muon_adamw' "
+                "to resume it"
+            )
+        if "nora" not in state_dict or "nadamw" not in state_dict:
+            raise ValueError(
+                "Checkpoint optimizer state is not compatible with nora_nadamw"
+            )
+
+        saved_nora = state_dict["nora"]
+        saved_nadamw = state_dict["nadamw"]
+        if (self.nora is None) != (saved_nora is None):
+            raise ValueError("Checkpoint Nora parameter groups do not match the model")
+        if (self.nadamw is None) != (saved_nadamw is None):
+            raise ValueError(
+                "Checkpoint NAdamW parameter groups do not match the model"
+            )
+        if self.nora is not None:
+            self.nora.load_state_dict(saved_nora)
+        if self.nadamw is not None:
+            self.nadamw.load_state_dict(saved_nadamw)
+        self._refresh_param_groups()

docs/README-zh-CN.md

@@ -107,7 +107,9 @@ nohup python scripts/tools/train.py \
     --batch_per_device=4 \
     --grad_accum_steps=8 \
     --warmup_ratio=0.05 \
+    --optimizer=nora_nadamw \
     --max_lr=1e-4 \
+    --nora_lr=5e-3 \
     --max_grad_norm=1.0 \
     --weight_decay=0.1 \
     --window_size=2048 \

docs/guides/params.md

@@ -25,21 +25,39 @@
 | Parameter | Description | Default |
 |-----------|-------------|---------|
 | `--warmup_ratio` | Fraction of total steps used for LR warmup | 0.05 |
-| `--max_lr` | Maximum learning rate (cosine decay after warmup) | 3e-4 |
+| `--max_lr` | NAdamW learning rate; schedulers scale every optimizer group proportionally | 3e-4 |
 | `--max_grad_norm` | Maximum gradient norm for clipping (None disables) | 1.0 |
 
-### Optimizer (MuonMix)
+### Optimizer
 
-Combined optimizer: matrix parameters via **Muon**, non-matrix via **AdamW** (`fused=True`).
+The default `nora_nadamw` optimizer sends internal `Linear.weight` matrices to
+**Nora** and embeddings, the LM head, norms, biases, LoRA factors, and fallback
+parameters to **NAdamW**. Parameters are classified by module role and identity,
+so tied embedding/head weights occur in exactly one group. Nora requires complete
+rows under DTensor sharding and rejects layouts sharded along the last dimension.
+
+| Parameter | Description | Default |
+|-----------|-------------|---------|
+| `--optimizer` | Built-in optimizer (`nora_nadamw`, `muon_adamw`) | `nora_nadamw` |
+| `--weight_decay` | NAdamW decay for eligible fallback parameters; known embeddings, heads, norms, biases, and LoRA factors use 0 | 0.1 |
+| `--nora_lr` | Nora learning rate | 5e-3 |
+| `--nora_beta` | Nora momentum-buffer EMA factor | 0.95 |
+| `--nora_momentum` | Nora Nesterov interpolation factor | 0.95 |
+| `--nora_weight_decay` | Nora matrix weight decay | 0.0 |
+
+`muon_adamw` preserves the previous MuonMix behavior and the following options:
 
 | Parameter | Description | Default |
 |-----------|-------------|---------|
-| `--weight_decay` | Weight decay (applied to Muon matrix params; non-matrix use 0) | 0.1 |
 | `--muon_momentum` | Muon momentum factor | 0.95 |
 | `--muon_nesterov` | Enable Nesterov momentum for Muon | True |
 | `--muon_ns_steps` | Newton-Schulz iteration steps for Muon | 5 |
 | `--muon_adjust_lr` | Muon LR adjustment strategy (`original`, `match_rms_adamw`) | `match_rms_adamw` |
 
+Optimizer identity and hyperparameters are saved in checkpoint metadata. Optimizer
+states are intentionally not interchangeable: resume older MuonMix checkpoints
+with `--optimizer=muon_adamw`.
+
 ### Data Loading
 
 | Parameter | Description | Default |
@@ -141,7 +159,9 @@ nohup python scripts/tools/train.py \
     --batch_per_device=4 \
     --grad_accum_steps=8 \
     --warmup_ratio=0.05 \
+    --optimizer=nora_nadamw \
     --max_lr=1e-4 \
+    --nora_lr=5e-3 \
     --max_grad_norm=1.0 \
     --weight_decay=0.1 \
     --window_size=2048 \

scripts/tools/train.py

@@ -1,115 +1,43 @@
 import os
 from collections.abc import Callable
 from functools import partial
-from typing import Any
 
 import click
 import torch
-from torch import Tensor, nn, optim
+from click.core import ParameterSource
+from torch import optim
 
 from astrai import setup_logging
 from astrai.config import AutoRegressiveLMConfig, TrainConfig
 from astrai.dataset import DatasetFactory, dpo_collate_fn, grpo_collate_fn
 from astrai.model import AutoRegressiveLM
 from astrai.model.components.decoder_block import DecoderBlock
+from astrai.optim import OptimizerFactory
 from astrai.trainer import SchedulerFactory, Trainer
 from astrai.trainer.rollout import BaseRewardModel
 
 
-class MuonMix(optim.Optimizer):
+def _merge_yaml_into_kwargs(
-    """Combined Muon (matrix) + AdamW (non-matrix) optimizer."""
+    config_path: str,
-
+    passed_kwargs: dict,
-    def __init__(
+    explicit_keys: set[str] | None = None,
-        self,
+) -> dict:
-        model: nn.Module,
+    """Merge Click defaults, YAML values, then explicit CLI values."""
-        lr: float = 3e-4,
-        weight_decay: float = 0.1,
-        momentum: float = 0.95,
-        nesterov: bool = True,
-        ns_steps: int = 5,
-        adjust_lr_fn: str = "match_rms_adamw",
-    ):
-        defaults = {
-            "lr": lr,
-            "weight_decay": weight_decay,
-            "momentum": momentum,
-            "nesterov": nesterov,
-            "ns_steps": ns_steps,
-            "adjust_lr_fn": adjust_lr_fn,
-        }
-        params = [p for p in model.parameters() if p.requires_grad]
-        super().__init__(params, defaults)
-
-        matrix_params: list[Tensor] = []
-        other_params: list[Tensor] = []
-        for name, param in model.named_parameters():
-            if not param.requires_grad:
-                continue
-            if (
-                param.dim() >= 2
-                and "norm" not in name
-                and "bias" not in name
-                and "embed" not in name
-                and "lm_head" not in name
-            ):
-                matrix_params.append(param)
-            else:
-                other_params.append(param)
-
-        self.muon = optim.Muon(
-            matrix_params,
-            lr=lr,
-            weight_decay=weight_decay,
-            momentum=momentum,
-            nesterov=nesterov,
-            ns_steps=ns_steps,
-            adjust_lr_fn=adjust_lr_fn,
-        )
-        self.adamw = optim.AdamW(
-            [{"params": other_params, "weight_decay": 0.0}],
-            lr=lr,
-            betas=(0.9, 0.95),
-            fused=True,
-        )
-
-        self.param_groups = [*self.muon.param_groups, *self.adamw.param_groups]
-
-    @torch.no_grad()
-    def step(self, closure=None):
-        self.muon.step(closure)
-        self.adamw.step(closure)
-
-    def zero_grad(self, set_to_none: bool = True):
-        self.muon.zero_grad(set_to_none)
-        self.adamw.zero_grad(set_to_none)
-
-    def state_dict(self) -> dict[str, Any]:
-        return {
-            "muon": self.muon.state_dict(),
-            "adamw": self.adamw.state_dict(),
-        }
-
-    def load_state_dict(self, state_dict: dict[str, Any]):
-        self.muon.load_state_dict(state_dict["muon"])
-        self.adamw.load_state_dict(state_dict["adamw"])
-        self.param_groups = [*self.muon.param_groups, *self.adamw.param_groups]
-
-
-def _merge_yaml_into_kwargs(config_path: str, passed_kwargs: dict) -> dict:
-    """Load YAML config, then override with explicit CLI kwargs (None excluded)."""
     import yaml
 
     with open(config_path) as f:
-        cfg = yaml.safe_load(f)
+        cfg = yaml.safe_load(f) or {}
 
-    merged = {}
+    merged = dict(passed_kwargs)
     for section in ("model", "data", "parallel", "training", "ckpt", "log"):
         if section in cfg:
             merged.update(cfg[section])
 
-    for key, value in passed_kwargs.items():
+    if explicit_keys is None:
-        if value is not None:
+        explicit_keys = set(passed_kwargs)
-            merged[key] = value
+    for key in explicit_keys:
+        if key in passed_kwargs:
+            merged[key] = passed_kwargs[key]
 
     return merged
 
@@ -117,6 +45,7 @@ def _merge_yaml_into_kwargs(config_path: str, passed_kwargs: dict) -> dict:
 _TRAIN_TYPE = ["seq", "sft", "dpo", "grpo", "online_grpo", "online_dpo"]
 _PARALLEL = ["none", "ddp", "fsdp"]
 _SCHEDULES = ["cosine", "sgdr", "wsd"]
+_OPTIMIZERS = OptimizerFactory.list_registered()
 _BACKENDS = ["nccl", "gloo"]
 _START_METHODS = ["spawn", "fork", "forkserver"]
 
@@ -162,10 +91,25 @@ _START_METHODS = ["spawn", "fork", "forkserver"]
     help="Fraction of total steps for LR warmup.",
 )
 @click.option("--max_lr", type=float, default=3e-4, help="Max learning rate.")
+@click.option(
+    "--optimizer",
+    type=click.Choice(_OPTIMIZERS),
+    default="nora_nadamw",
+    help="Built-in optimizer.",
+)
 @click.option(
     "--max_grad_norm", type=float, default=1.0, help="Max gradient norm for clipping."
 )
-@click.option("--weight_decay", type=float, default=0.1, help="Weight decay.")
+@click.option(
+    "--weight_decay",
+    type=float,
+    default=0.1,
+    help="Weight decay for eligible optimizer parameters.",
+)
+@click.option("--nora_lr", type=float, default=5e-3, help="Nora learning rate.")
+@click.option("--nora_beta", type=float, default=0.95, help="Nora EMA factor.")
+@click.option("--nora_momentum", type=float, default=0.95, help="Nora update momentum.")
+@click.option("--nora_weight_decay", type=float, default=0.0, help="Nora weight decay.")
 @click.option("--muon_momentum", type=float, default=0.95, help="Muon momentum factor.")
 @click.option("--muon_nesterov/--no-muon_nesterov", default=True, help="Muon Nesterov.")
 @click.option("--muon_ns_steps", type=int, default=5, help="Muon Newton-Schulz steps.")
@@ -283,8 +227,14 @@ _START_METHODS = ["spawn", "fork", "forkserver"]
 @click.pass_context
 def train_command(ctx, config_path, dry_run, metrics, **kwargs):
     """Start model training (pretrain / SFT / DPO / GRPO)."""
+    kwargs["metrics"] = metrics
     if config_path:
-        kwargs = _merge_yaml_into_kwargs(config_path, kwargs)
+        explicit_keys = {
+            key
+            for key in kwargs
+            if ctx.get_parameter_source(key) is ParameterSource.COMMANDLINE
+        }
+        kwargs = _merge_yaml_into_kwargs(config_path, kwargs, explicit_keys)
 
     required = ["train_type", "data_root_path", "param_path"]
     missing = [k for k in required if kwargs.get(k) is None]
@@ -295,7 +245,7 @@ def train_command(ctx, config_path, dry_run, metrics, **kwargs):
         )
 
     # Convert tuple back to list
-    kwargs["metrics"] = list(metrics)
+    kwargs["metrics"] = list(kwargs["metrics"])
     # Remove tp_size (not yet wired)
     kwargs.pop("tp_size", None)
 
@@ -317,6 +267,7 @@ def _print_dry_run(kwargs: dict) -> None:
         ("Epochs", str(kwargs.get("n_epoch", 1))),
         ("Batch/device", str(kwargs.get("batch_per_device", 1))),
         ("Grad accum", str(kwargs.get("grad_accum_steps", 1))),
+        ("Optimizer", str(kwargs.get("optimizer", "nora_nadamw"))),
         ("Max LR", str(kwargs.get("max_lr", "?"))),
         ("Schedule", str(kwargs.get("schedule_type", "cosine"))),
         ("Warmup ratio", str(kwargs.get("warmup_ratio", 0.05))),
@@ -336,8 +287,10 @@ def create_model(config):
     return AutoRegressiveLM(config).to(dtype=torch.bfloat16)
 
 
-def create_optimizer(model, **kwargs) -> MuonMix:
+def create_optimizer(
-    return MuonMix(model, **kwargs)
+    model, optimizer_name: str = "nora_nadamw", **kwargs
+) -> optim.Optimizer:
+    return OptimizerFactory.create(optimizer_name, model, **kwargs)
 
 
 def create_scheduler(
@@ -459,15 +412,51 @@ def train(
         tokenizer_path=param_path,
     )
 
+    optimizer_name = kwargs.pop("optimizer", "nora_nadamw")
+    optimizer_kwargs = {
+        "lr": kwargs.pop("max_lr"),
+        "weight_decay": kwargs.pop("weight_decay"),
+        "nora_lr": kwargs.pop("nora_lr", 5e-3),
+        "nora_beta": kwargs.pop("nora_beta", 0.95),
+        "nora_momentum": kwargs.pop("nora_momentum", 0.95),
+        "nora_weight_decay": kwargs.pop("nora_weight_decay", 0.0),
+        "momentum": kwargs.pop("muon_momentum", 0.95),
+        "nesterov": kwargs.pop("muon_nesterov", True),
+        "ns_steps": kwargs.pop("muon_ns_steps", 5),
+        "adjust_lr_fn": kwargs.pop("muon_adjust_lr", "match_rms_adamw"),
+    }
     optimizer_fn = partial(
         create_optimizer,
-        lr=kwargs.pop("max_lr"),
+        optimizer_name=optimizer_name,
-        weight_decay=kwargs.pop("weight_decay"),
+        **optimizer_kwargs,
-        momentum=kwargs.pop("muon_momentum"),
-        nesterov=kwargs.pop("muon_nesterov"),
-        ns_steps=kwargs.pop("muon_ns_steps"),
-        adjust_lr_fn=kwargs.pop("muon_adjust_lr"),
     )
+    if optimizer_name == "nora_nadamw":
+        optimizer_hyperparameters = {
+            key: optimizer_kwargs[key]
+            for key in (
+                "lr",
+                "weight_decay",
+                "nora_lr",
+                "nora_beta",
+                "nora_momentum",
+                "nora_weight_decay",
+            )
+        }
+        optimizer_hyperparameters.update(
+            {"nadamw_betas": [0.9, 0.999], "nadamw_eps": 1e-8, "nora_eps": 1e-10}
+        )
+    else:
+        optimizer_hyperparameters = {
+            key: optimizer_kwargs[key]
+            for key in (
+                "lr",
+                "weight_decay",
+                "momentum",
+                "nesterov",
+                "ns_steps",
+                "adjust_lr_fn",
+            )
+        }
 
     total_steps = compute_total_steps(
         len(dataset), n_epoch, batch_per_device, nprocs, grad_accum_steps
@@ -516,6 +505,8 @@ def train(
         dataset=dataset,
         optimizer_fn=optimizer_fn,
         scheduler_fn=scheduler_fn,
+        optimizer_name=optimizer_name,
+        optimizer_hyperparameters=optimizer_hyperparameters,
         ckpt_dir=ckpt_dir,
         n_epoch=n_epoch,
         batch_per_device=batch_per_device,

tests/optim/test_nora_nadamw.py

@@ -0,0 +1,253 @@
+import math
+from copy import deepcopy
+
+import pytest
+import torch
+from torch.utils.data import TensorDataset
+
+from astrai.config import TrainConfig
+from astrai.model import AutoRegressiveLM
+from astrai.model.components.linear import Linear
+from astrai.model.components.lora import LoRALinear, inject_lora
+from astrai.optim import (
+    NAdamW,
+    Nora,
+    NoraNAdamW,
+    OptimizerFactory,
+    nora_lr_scale,
+    partition_optimizer_parameters,
+)
+from astrai.trainer.schedule import SchedulerFactory
+from tests.helpers import make_tiny_config
+
+
+def _set_constant_grads(model, value):
+    for param in model.parameters():
+        if param.requires_grad:
+            param.grad = torch.full_like(param, value)
+
+
+def test_nora_one_step_matches_row_geometry():
+    param = torch.nn.Parameter(torch.tensor([[3.0, 4.0], [0.0, 2.0]]))
+    grad = torch.tensor([[4.0, -3.0], [1.0, 1.0]])
+    param.grad = grad.clone()
+
+    optimizer = Nora([param], lr=0.1, beta=0.0, momentum=0.0)
+    optimizer.step()
+
+    theta_hat = torch.tensor([[0.6, 0.8], [0.0, 1.0]])
+    tangent = grad - (grad * theta_hat).sum(dim=-1, keepdim=True) * theta_hat
+    direction = tangent / tangent.norm(dim=-1, keepdim=True).clamp(min=1e-10)
+    expected = torch.tensor([[3.0, 4.0], [0.0, 2.0]]) - 0.1 * direction
+    torch.testing.assert_close(param, expected)
+
+
+def test_nora_handles_zero_and_pure_radial_rows():
+    param = torch.nn.Parameter(torch.tensor([[0.0, 0.0], [3.0, 4.0]]))
+    param.grad = torch.tensor([[3.0, 4.0], [6.0, 8.0]])
+
+    optimizer = Nora([param], lr=0.1, beta=0.0, momentum=0.0)
+    optimizer.step()
+
+    torch.testing.assert_close(param[0], torch.tensor([-0.06, -0.08]))
+    torch.testing.assert_close(param[1], torch.tensor([3.0, 4.0]), atol=1e-6, rtol=0)
+
+
+def test_nora_lr_scale_only_increases_tall_matrices():
+    assert nora_lr_scale(0.1, torch.Size([4, 2])) == pytest.approx(0.1 * math.sqrt(2.0))
+    assert nora_lr_scale(0.1, torch.Size([2, 4])) == pytest.approx(0.1)
+
+
+def test_nadamw_one_step_matches_reference_formula():
+    param = torch.nn.Parameter(torch.tensor([1.0, -2.0]))
+    grad = torch.tensor([0.5, -0.25])
+    param.grad = grad.clone()
+    lr = 0.1
+    beta1, beta2 = 0.9, 0.999
+    eps = 1e-8
+
+    optimizer = NAdamW([param], lr=lr, betas=(beta1, beta2), eps=eps, weight_decay=0.2)
+    optimizer.step()
+
+    m = (1 - beta1) * grad
+    v = (1 - beta2) * grad.square()
+    m_hat = (beta1 * m + (1 - beta1) * grad) / (1 - beta1)
+    v_hat = v / (1 - beta2)
+    expected = torch.tensor([1.0, -2.0]) * (1 - lr * 0.2)
+    expected.add_(m_hat / (v_hat.sqrt() + eps), alpha=-lr)
+    torch.testing.assert_close(param, expected)
+
+
+@pytest.mark.parametrize("tie_word_embeddings", [False, True])
+def test_parameter_partition_is_complete_disjoint_and_role_based(
+    tie_word_embeddings,
+):
+    model = AutoRegressiveLM(make_tiny_config(tie_word_embeddings=tie_word_embeddings))
+    inject_lora(model, r=2, alpha=4, target_modules={"q_proj"})
+
+    groups = partition_optimizer_parameters(model)
+    all_grouped = [*groups.nora, *groups.nadamw_decay, *groups.nadamw_no_decay]
+    trainable = [param for param in model.parameters() if param.requires_grad]
+
+    assert len({id(param) for param in all_grouped}) == len(all_grouped)
+    assert {id(param) for param in all_grouped} == {id(param) for param in trainable}
+    assert id(model.embed_tokens.weight) in {id(p) for p in groups.nadamw_no_decay}
+    assert id(model.lm_head.weight) in {id(p) for p in groups.nadamw_no_decay}
+
+    nora_ids = {id(param) for param in groups.nora}
+    no_decay_ids = {id(param) for param in groups.nadamw_no_decay}
+    for name, module in model.named_modules():
+        if isinstance(module, LoRALinear):
+            assert id(module.lora_A) in no_decay_ids
+            assert id(module.lora_B) in no_decay_ids
+        elif isinstance(module, Linear) and name != "lm_head":
+            if module.weight.requires_grad:
+                assert id(module.weight) in nora_ids
+
+
+@pytest.mark.parametrize(
+    "model_overrides",
+    [
+        {"attn_type": "gqa", "ffn_type": "mlp"},
+        {
+            "attn_type": "gqa",
+            "ffn_type": "moe",
+            "n_routed_experts": 2,
+            "n_shared_experts": 1,
+            "n_activated_experts": 1,
+            "topk_method": "greedy",
+        },
+        {
+            "attn_type": "mla",
+            "ffn_type": "mlp",
+            "kv_lora_rank": 4,
+            "qk_nope_head_dim": 2,
+            "qk_rope_head_dim": 2,
+        },
+        {
+            "attn_type": "mla",
+            "ffn_type": "moe",
+            "kv_lora_rank": 4,
+            "qk_nope_head_dim": 2,
+            "qk_rope_head_dim": 2,
+            "n_routed_experts": 2,
+            "n_shared_experts": 1,
+            "n_activated_experts": 1,
+            "topk_method": "greedy",
+        },
+    ],
+)
+def test_parameter_partition_covers_all_model_structures(model_overrides):
+    model = AutoRegressiveLM(make_tiny_config(**model_overrides))
+    groups = partition_optimizer_parameters(model)
+
+    grouped = [*groups.nora, *groups.nadamw_decay, *groups.nadamw_no_decay]
+    trainable = [param for param in model.parameters() if param.requires_grad]
+    assert {id(param) for param in grouped} == {id(param) for param in trainable}
+    assert len(grouped) == len({id(param) for param in grouped})
+
+
+def test_factory_registers_nora_default_and_legacy_muon():
+    assert OptimizerFactory.list_registered() == ["muon_adamw", "nora_nadamw"]
+    model = AutoRegressiveLM(make_tiny_config())
+    optimizer = OptimizerFactory.create("nora_nadamw", model, lr=3e-4)
+    assert isinstance(optimizer, NoraNAdamW)
+
+
+def test_scheduler_preserves_nora_to_nadamw_lr_ratio():
+    model = AutoRegressiveLM(make_tiny_config())
+    optimizer = NoraNAdamW(model, lr=3e-4, nora_lr=5e-3)
+    scheduler = SchedulerFactory.create(
+        "cosine", optimizer, warmup_steps=2, lr_decay_steps=2, min_rate=0.1
+    )
+
+    initial_ratio = optimizer.param_groups[0]["lr"] / optimizer.param_groups[-1]["lr"]
+    _set_constant_grads(model, 0.1)
+    optimizer.step()
+    scheduler.step()
+    stepped_ratio = optimizer.param_groups[0]["lr"] / optimizer.param_groups[-1]["lr"]
+
+    assert initial_ratio == pytest.approx(5e-3 / 3e-4)
+    assert stepped_ratio == pytest.approx(initial_ratio)
+
+
+def test_optimizer_and_scheduler_resume_matches_uninterrupted_step():
+    torch.manual_seed(7)
+    model_a = AutoRegressiveLM(make_tiny_config())
+    optimizer_a = NoraNAdamW(model_a, lr=3e-4, nora_lr=5e-3)
+    scheduler_a = SchedulerFactory.create(
+        "cosine", optimizer_a, warmup_steps=2, lr_decay_steps=4, min_rate=0.1
+    )
+
+    _set_constant_grads(model_a, 0.125)
+    optimizer_a.step()
+    scheduler_a.step()
+    model_state = {key: value.clone() for key, value in model_a.state_dict().items()}
+    optimizer_state = deepcopy(optimizer_a.state_dict())
+    scheduler_state = deepcopy(scheduler_a.state_dict())
+
+    model_b = AutoRegressiveLM(make_tiny_config())
+    model_b.load_state_dict(model_state)
+    optimizer_b = NoraNAdamW(model_b, lr=3e-4, nora_lr=5e-3)
+    scheduler_b = SchedulerFactory.create(
+        "cosine", optimizer_b, warmup_steps=2, lr_decay_steps=4, min_rate=0.1
+    )
+    optimizer_b.load_state_dict(optimizer_state)
+    scheduler_b.load_state_dict(scheduler_state)
+
+    _set_constant_grads(model_a, -0.25)
+    _set_constant_grads(model_b, -0.25)
+    optimizer_a.step()
+    optimizer_b.step()
+    scheduler_a.step()
+    scheduler_b.step()
+
+    for param_a, param_b in zip(model_a.parameters(), model_b.parameters()):
+        torch.testing.assert_close(param_a, param_b)
+    assert scheduler_a.get_last_lr() == pytest.approx(scheduler_b.get_last_lr())
+
+
+def test_train_config_serializes_optimizer_metadata():
+    config = TrainConfig(
+        model_fn=lambda: torch.nn.Linear(2, 2),
+        strategy="seq",
+        dataset=TensorDataset(torch.zeros(1, 2)),
+        optimizer_fn=lambda model: torch.optim.AdamW(model.parameters()),
+        scheduler_fn=lambda optimizer: torch.optim.lr_scheduler.LambdaLR(
+            optimizer, lambda _: 1.0
+        ),
+        optimizer_name="nora_nadamw",
+        optimizer_hyperparameters={"lr": 3e-4, "nora_lr": 5e-3},
+    )
+
+    metadata = config.to_dict()
+
+    assert metadata["optimizer_name"] == "nora_nadamw"
+    assert metadata["optimizer_hyperparameters"] == {
+        "lr": 3e-4,
+        "nora_lr": 5e-3,
+    }
+
+
+def test_nora_nadamw_rejects_legacy_muon_state():
+    model = AutoRegressiveLM(make_tiny_config())
+    optimizer = NoraNAdamW(model)
+
+    with pytest.raises(ValueError, match="muon_adamw"):
+        optimizer.load_state_dict({"muon": {}, "adamw": {}})
+
+
+def test_combined_optimizer_runs_closure_once():
+    model = AutoRegressiveLM(make_tiny_config())
+    optimizer = NoraNAdamW(model)
+    calls = 0
+
+    def closure():
+        nonlocal calls
+        calls += 1
+        return torch.tensor(1.0, requires_grad=True)
+
+    loss = optimizer.step(closure)
+
+    assert calls == 1
+    assert loss.item() == 1.0

tests/optim/test_optimizer_distributed.py

@@ -0,0 +1,65 @@
+import pytest
+import torch
+import torch.distributed as dist
+from torch.distributed.fsdp import fully_shard
+from torch.distributed.tensor import DTensor, Shard
+from torch.nn.parallel import DistributedDataParallel as DDP
+
+from astrai.model import AutoRegressiveLM
+from astrai.optim import NoraNAdamW
+from astrai.parallel.setup import find_free_port
+from tests.helpers import make_tiny_config
+
+pytestmark = pytest.mark.skipif(
+    torch.cuda.device_count() < 1, reason="CUDA device required"
+)
+
+
+def _assign_grads_and_step(model):
+    optimizer = NoraNAdamW(model)
+    for param in model.parameters():
+        if param.requires_grad:
+            param.grad = torch.ones_like(param)
+    optimizer.step()
+    return optimizer
+
+
+def test_nora_nadamw_steps_after_ddp_and_fsdp2_wrapping():
+    torch.cuda.set_device(0)
+    dist.init_process_group(
+        "nccl",
+        rank=0,
+        world_size=1,
+        init_method=f"tcp://127.0.0.1:{find_free_port()}",
+    )
+    try:
+        ddp_model = AutoRegressiveLM(make_tiny_config()).to(
+            device="cuda", dtype=torch.bfloat16
+        )
+        ddp_model = DDP(ddp_model, device_ids=[0], output_device=0)
+        ddp_optimizer = _assign_grads_and_step(ddp_model)
+        assert ddp_optimizer.state_dict()["nora"]["state"]
+
+        fsdp_model = AutoRegressiveLM(make_tiny_config()).to(
+            device="cuda", dtype=torch.bfloat16
+        )
+        for child in fsdp_model.children():
+            if isinstance(child, torch.nn.ModuleList):
+                for submodule in child:
+                    fully_shard(submodule, reshard_after_forward=False)
+            else:
+                fully_shard(child, reshard_after_forward=False)
+
+        fsdp_optimizer = _assign_grads_and_step(fsdp_model)
+        nora_params = fsdp_optimizer.nora.param_groups[0]["params"]
+        assert nora_params
+        assert all(isinstance(param, DTensor) for param in nora_params)
+        assert all(
+            all(
+                not isinstance(placement, Shard) or placement.dim == 0
+                for placement in param.placements
+            )
+            for param in nora_params
+        )
+    finally:
+        dist.destroy_process_group()

tests/test_train_cli.py

@@ -0,0 +1,62 @@
+import re
+
+from click.testing import CliRunner
+
+from scripts.tools.train import _merge_yaml_into_kwargs, train_command
+
+
+def test_yaml_overrides_click_defaults_but_not_explicit_cli(tmp_path):
+    config_path = tmp_path / "train.yaml"
+    config_path.write_text(
+        "training:\n"
+        "  optimizer: nora_nadamw\n"
+        "  max_lr: 0.0002\n"
+        "  nora_lr: 0.004\n"
+        "  batch_per_device: 8\n",
+        encoding="utf-8",
+    )
+    click_values = {
+        "optimizer": "nora_nadamw",
+        "max_lr": 3e-4,
+        "nora_lr": 5e-3,
+        "batch_per_device": 16,
+    }
+
+    merged = _merge_yaml_into_kwargs(
+        str(config_path), click_values, explicit_keys={"batch_per_device"}
+    )
+
+    assert merged["max_lr"] == 2e-4
+    assert merged["nora_lr"] == 4e-3
+    assert merged["batch_per_device"] == 16
+
+
+def test_train_dry_run_uses_yaml_then_explicit_cli(tmp_path):
+    data_path = tmp_path / "data"
+    model_path = tmp_path / "model"
+    data_path.mkdir()
+    model_path.mkdir()
+    config_path = tmp_path / "train.yaml"
+    config_path.write_text(
+        "data:\n"
+        f"  data_root_path: {data_path}\n"
+        "model:\n"
+        f"  param_path: {model_path}\n"
+        "training:\n"
+        "  train_type: seq\n"
+        "  optimizer: nora_nadamw\n"
+        "  max_lr: 0.0002\n"
+        "  nora_lr: 0.004\n"
+        "  batch_per_device: 8\n",
+        encoding="utf-8",
+    )
+
+    result = CliRunner().invoke(
+        train_command,
+        ["--config", str(config_path), "--dry-run", "--batch_per_device", "16"],
+    )
+
+    assert result.exit_code == 0, result.output
+    assert re.search(r"Optimizer\s+: nora_nadamw", result.output)
+    assert re.search(r"Batch/device\s+: 16", result.output)
+    assert re.search(r"Max LR\s+: 0.0002", result.output)