chore: bump version to 1.3.12

- ManoAdamW now uses torch.optim.AdamW(fused=True, betas=(0.9, 0.95)) matching MuonAdamW, eliminating a confounding variable in optimizer comparison experiments
- only NoraNAdamW retains NAdamW, which is correct per the Nora paper design

astrai/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "1.3.11"
+__version__ = "1.3.12"
 __author__ = "ViperEkura"
 
 import logging

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/extension/attention_backend.py

@@ -30,7 +30,6 @@ Layout convention: all q/k/v are ``[batch, seq_len, n_heads, head_dim]``
 
 import contextvars
 import enum
-import math
 from abc import ABC, abstractmethod
 from contextlib import contextmanager
 from typing import Optional, Union

astrai/optim/__init__.py

@@ -0,0 +1,38 @@
+"""Optimizer implementations and factory registration."""
+
+from astrai.optim.composite import (
+    OptimizerFactory,
+    composite_state_dict,
+    composite_step,
+    composite_zero_grad,
+    refresh_param_groups,
+)
+from astrai.optim.mano_adamw import Mano, ManoAdamW
+from astrai.optim.muon_adamw import MuonAdamW
+from astrai.optim.nora_nadamw import (
+    NAdamW,
+    Nora,
+    NoraNAdamW,
+    OptimizerParameterGroups,
+    nora_direction,
+    nora_lr_scale,
+    partition_optimizer_parameters,
+)
+
+__all__ = [
+    "Mano",
+    "ManoAdamW",
+    "MuonAdamW",
+    "NAdamW",
+    "Nora",
+    "NoraNAdamW",
+    "OptimizerFactory",
+    "OptimizerParameterGroups",
+    "composite_state_dict",
+    "composite_step",
+    "composite_zero_grad",
+    "nora_direction",
+    "nora_lr_scale",
+    "partition_optimizer_parameters",
+    "refresh_param_groups",
+]

astrai/optim/composite.py

@@ -0,0 +1,71 @@
+"""Shared infrastructure for the optim package.
+
+This module hosts two things:
+
+* ``OptimizerFactory`` — the registry for built-in optimizers. Defining it
+  here (rather than in ``__init__.py``) lets each optimizer module import it
+  and register itself with a decorator, avoiding circular imports.
+* Composite-optimizer helpers — ``step``/``zero_grad``/``state_dict``/
+  ``param_groups`` delegation shared by every optimizer that routes different
+  parameter groups through distinct sub-optimizers.
+"""
+
+from typing import Any
+
+import torch
+from torch.optim import Optimizer
+
+from astrai.factory import BaseFactory
+
+
+class OptimizerFactory(BaseFactory[Optimizer]):
+    """Factory for built-in training optimizers."""
+
+
+def composite_step(
+    sub_optimizers: list[Optimizer],
+    closure=None,
+) -> torch.Tensor | None:
+    """Run ``step`` on every sub-optimizer, invoking the closure once.
+
+    The closure (if given) is executed inside ``torch.enable_grad`` exactly
+    once before any sub-optimizer steps, matching the contract of a single
+    ``Optimizer.step``. Sub-optimizers receive ``None`` so they do not
+    re-execute it.
+    """
+    loss = None
+    if closure is not None:
+        with torch.enable_grad():
+            loss = closure()
+    for sub in sub_optimizers:
+        sub.step()
+    return loss
+
+
+def composite_zero_grad(
+    sub_optimizers: list[Optimizer],
+    set_to_none: bool = True,
+) -> None:
+    for sub in sub_optimizers:
+        sub.zero_grad(set_to_none=set_to_none)
+
+
+def composite_state_dict(
+    named_sub_optimizers: dict[str, Optimizer | None],
+) -> dict[str, Any]:
+    """Serialize sub-optimizers, preserving ``None`` slots."""
+    return {
+        name: sub.state_dict() if sub is not None else None
+        for name, sub in named_sub_optimizers.items()
+    }
+
+
+def refresh_param_groups(
+    sub_optimizers: list[Optimizer],
+) -> list[dict]:
+    """Concatenate param_groups from every non-None sub-optimizer."""
+    groups: list[dict] = []
+    for sub in sub_optimizers:
+        if sub is not None:
+            groups.extend(sub.param_groups)
+    return groups

astrai/optim/mano_adamw.py

@@ -0,0 +1,214 @@
+"""Mano manifold optimizer combined with AdamW.
+
+Mano projects the momentum onto the tangent space of the Oblique manifold
+(axis-wise tangent projection) and normalizes it, replacing the expensive
+Newton-Schulz iteration in Muon with a cheaper manifold normalization.
+
+Reference: https://arxiv.org/abs/2601.23000
+"""
+
+import math
+
+import torch
+from torch import nn, optim
+from torch.optim import Optimizer
+
+from astrai.optim.composite import (
+    OptimizerFactory,
+    composite_state_dict,
+    composite_step,
+    composite_zero_grad,
+    refresh_param_groups,
+)
+from astrai.optim.nora_nadamw import partition_optimizer_parameters
+
+
+class Mano(Optimizer):
+    """Manifold Normalized Optimizer for two-dimensional matrices.
+
+    Each step alternates the projection axis (dim 0 / dim 1) to restrike the
+    manifold along both rows and columns. The tangent momentum is computed
+    without normalizing the parameter itself (v2 simplification) and the
+    epsilon is added (not clamped) to the norm denominator.
+    """
+
+    def __init__(
+        self,
+        params,
+        lr: float = 1e-3,
+        weight_decay: float = 0.1,
+        momentum: float = 0.95,
+        nesterov: bool = True,
+        eps: float = 1e-8,
+    ):
+        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 eps <= 0:
+            raise ValueError(f"Invalid epsilon: {eps}")
+
+        defaults = {
+            "lr": lr,
+            "weight_decay": weight_decay,
+            "momentum": momentum,
+            "nesterov": nesterov,
+            "eps": eps,
+            "steps": 0,
+        }
+        super().__init__(params, defaults)
+        for group in self.param_groups:
+            for param in group["params"]:
+                if param.ndim != 2:
+                    raise ValueError(
+                        f"Mano only supports 2D matrices, got shape {tuple(param.shape)}"
+                    )
+
+    @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"]
+            nesterov = group["nesterov"]
+            eps = group["eps"]
+            dim = int(group["steps"] % 2)
+
+            for param in group["params"]:
+                if param.grad is None:
+                    continue
+                if param.grad.is_sparse:
+                    raise RuntimeError("Mano 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.mul_(momentum).add_(grad)
+                update = (
+                    grad.add(momentum_buffer, alpha=momentum)
+                    if nesterov
+                    else momentum_buffer
+                )
+
+                tangent = update - (
+                    torch.sum(update * param.data, dim=dim, keepdim=True) * param.data
+                )
+                direction = tangent / (
+                    torch.norm(tangent, p=2, dim=dim, keepdim=True) + eps
+                )
+
+                if weight_decay != 0:
+                    param.mul_(1 - lr * weight_decay)
+                adjusted_lr = lr * 0.2 * math.sqrt(direction.shape[dim])
+                param.add_(direction, alpha=-adjusted_lr)
+                state["momentum_buffer"] = momentum_buffer
+
+            group["steps"] += 1
+
+        return loss
+
+
+@OptimizerFactory.register("mano_adamw")
+class ManoAdamW(Optimizer):
+    """Mano for internal linear weights and AdamW for remaining parameters."""
+
+    optimizer_name = "mano_adamw"
+
+    def __init__(
+        self,
+        model: nn.Module,
+        lr: float = 3e-4,
+        weight_decay: float = 0.1,
+        momentum: float = 0.95,
+        nesterov: bool = True,
+    ):
+        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.mano = (
+            Mano(
+                groups.nora,
+                lr=lr,
+                weight_decay=weight_decay,
+                momentum=momentum,
+                nesterov=nesterov,
+            )
+            if groups.nora
+            else None
+        )
+
+        adamw_groups = []
+        if groups.nadamw_decay:
+            adamw_groups.append(
+                {"params": groups.nadamw_decay, "weight_decay": weight_decay}
+            )
+        if groups.nadamw_no_decay:
+            adamw_groups.append({"params": groups.nadamw_no_decay, "weight_decay": 0.0})
+        self.adamw = (
+            optim.AdamW(
+                adamw_groups,
+                lr=lr,
+                betas=(0.9, 0.95),
+                fused=True,
+            )
+            if adamw_groups
+            else None
+        )
+        self.param_groups = refresh_param_groups([self.mano, self.adamw])
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        return composite_step(
+            [opt for opt in (self.mano, self.adamw) if opt is not None],
+            closure,
+        )
+
+    def zero_grad(self, set_to_none: bool = True):
+        composite_zero_grad(
+            [opt for opt in (self.mano, self.adamw) if opt is not None],
+            set_to_none,
+        )
+
+    def state_dict(self) -> dict:
+        return composite_state_dict({"mano": self.mano, "adamw": self.adamw})
+
+    def load_state_dict(self, state_dict: dict):
+        if "muon" in state_dict or "nora" in state_dict:
+            raise ValueError(
+                "Checkpoint uses a different optimizer; select the matching "
+                "--optimizer to resume it"
+            )
+        if "mano" not in state_dict or "adamw" not in state_dict:
+            raise ValueError(
+                "Checkpoint optimizer state is not compatible with mano_adamw"
+            )
+
+        saved_mano = state_dict["mano"]
+        saved_adamw = state_dict["adamw"]
+        if (self.mano is None) != (saved_mano is None):
+            raise ValueError("Checkpoint Mano parameter groups do not match the model")
+        if (self.adamw is None) != (saved_adamw is None):
+            raise ValueError("Checkpoint AdamW parameter groups do not match the model")
+        if self.mano is not None:
+            self.mano.load_state_dict(saved_mano)
+        if self.adamw is not None:
+            self.adamw.load_state_dict(saved_adamw)
+        self.param_groups = refresh_param_groups([self.mano, self.adamw])

astrai/optim/muon_adamw.py

@@ -0,0 +1,95 @@
+"""Legacy Muon + AdamW combined optimizer."""
+
+from typing import Any
+
+import torch
+from torch import Tensor, nn, optim
+
+from astrai.optim.composite import (
+    OptimizerFactory,
+    composite_state_dict,
+    composite_step,
+    composite_zero_grad,
+    refresh_param_groups,
+)
+
+
+@OptimizerFactory.register("muon_adamw")
+class MuonAdamW(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 = refresh_param_groups([self.muon, self.adamw])
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        return composite_step([self.muon, self.adamw], closure)
+
+    def zero_grad(self, set_to_none: bool = True):
+        composite_zero_grad([self.muon, self.adamw], set_to_none)
+
+    def state_dict(self) -> dict[str, Any]:
+        return composite_state_dict({"muon": self.muon, "adamw": self.adamw})
+
+    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 = refresh_param_groups([self.muon, self.adamw])

astrai/optim/nora_nadamw.py

@@ -0,0 +1,372 @@
+"""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
+from astrai.optim.composite import (
+    OptimizerFactory,
+    composite_state_dict,
+    composite_step,
+    composite_zero_grad,
+    refresh_param_groups,
+)
+
+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)
+
+
+@OptimizerFactory.register("nora_nadamw")
+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.param_groups = refresh_param_groups([self.nora, self.nadamw])
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        return composite_step(
+            [opt for opt in (self.nora, self.nadamw) if opt is not None],
+            closure,
+        )
+
+    def zero_grad(self, set_to_none: bool = True):
+        composite_zero_grad(
+            [opt for opt in (self.nora, self.nadamw) if opt is not None],
+            set_to_none,
+        )
+
+    def state_dict(self) -> dict[str, Any]:
+        return composite_state_dict({"nora": self.nora, "nadamw": self.nadamw})
+
+    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.param_groups = refresh_param_groups([self.nora, self.nadamw])

astrai/trainer/metric_util.py

@@ -22,6 +22,51 @@ def grad_norm(model: nn.Module, per_param: bool = False) -> float | Dict[str, fl
     return total_sq.sqrt().item()
 
 
+class GradSNRTracker:
+    """Track gradient signal-to-noise ratio via EMA of first/second moments.
+
+    SNR = E[g]^2 / Var(g) = E[g]^2 / (E[g^2] - E[g]^2)
+
+    The tracker accumulates per-parameter EMA moments across optimizer steps.
+    Call ``update`` after backward (before ``optimizer.step``) and read
+    ``snr`` to get the aggregate SNR across all parameters.
+    """
+
+    def __init__(self, beta: float = 0.999, eps: float = 1e-8):
+        self.beta = beta
+        self.eps = eps
+        self._first: Dict[int, torch.Tensor] = {}
+        self._second: Dict[int, torch.Tensor] = {}
+
+    @torch.no_grad()
+    def update(self, model: nn.Module) -> None:
+        beta = self.beta
+        for param in model.parameters():
+            if param.grad is None:
+                continue
+            pid = id(param)
+            g = param.grad.detach()
+            if pid not in self._first:
+                self._first[pid] = g.clone()
+                self._second[pid] = g.pow(2).clone()
+            else:
+                self._first[pid].mul_(beta).add_(g, alpha=1 - beta)
+                self._second[pid].mul_(beta).addcmul_(g, g, value=1 - beta)
+
+    @property
+    def snr(self) -> float:
+        if not self._first:
+            return 0.0
+        total_signal = 0.0
+        total_noise = 0.0
+        for m, v in zip(self._first.values(), self._second.values()):
+            signal = m.pow(2).sum().item()
+            noise = (v - m.pow(2)).clamp(min=0).sum().item()
+            total_signal += signal
+            total_noise += noise
+        return total_signal / (total_noise + self.eps)
+
+
 def ctx_get_loss(ctx):
     return ctx.loss
 
@@ -36,3 +81,10 @@ def ctx_get_val_loss(ctx):
 
 def ctx_get_grad_norm(ctx):
     return ctx.grad_norm
+
+
+def ctx_get_grad_snr(ctx):
+    tracker = getattr(ctx, "grad_snr_tracker", None)
+    if tracker is None:
+        return None
+    return tracker.snr

astrai/trainer/train_callback.py

@@ -18,6 +18,7 @@ from astrai.parallel.setup import get_current_device
 from astrai.serialization import Checkpoint
 from astrai.trainer.metric_util import (
     ctx_get_grad_norm,
+    ctx_get_grad_snr,
     ctx_get_loss,
     ctx_get_lr,
     ctx_get_val_loss,
@@ -255,6 +256,7 @@ class MetricCallback(TrainCallback):
             "lr": ctx_get_lr,
             "val_loss": ctx_get_val_loss,
             "grad_norm": ctx_get_grad_norm,
+            "grad_snr": ctx_get_grad_snr,
         }
 
     def _metrics(self, context: TrainContext, names):
@@ -312,6 +314,8 @@ class MetricCallback(TrainCallback):
                 f.write(json.dumps(log) + "\n")
 
     def on_optimizer_step(self, context):
+        context.grad_snr_tracker.update(context.model)
+
         if (
             context.val_dataloader is not None
             and self.val_step > 0

astrai/trainer/train_context.py

@@ -17,6 +17,7 @@ from astrai.parallel.setup import get_current_device, get_rank, get_world_size
 from astrai.protocols import OptimizerProtocol, SchedulerProtocol
 from astrai.serialization import Checkpoint, load_json
 from astrai.tokenize import AutoTokenizer
+from astrai.trainer.metric_util import GradSNRTracker
 from astrai.trainer.rollout import RolloutGenerator, RolloutRunner
 from astrai.trainer.strategy import BaseStrategy, StrategyFactory
 
@@ -38,6 +39,7 @@ class TrainContext:
     consumed_samples: int = field(default=0)
     loss: float = field(default=0.0)
     grad_norm: Optional[float] = field(default=None)
+    grad_snr_tracker: GradSNRTracker = field(default_factory=GradSNRTracker)
     val_dataloader: Optional[DataLoader] = field(default=None)
     val_loss: Optional[float] = field(default=None)
 

csrc/kernels/attn_decode_split_kv_mma.cuh

@@ -76,26 +76,17 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
         cp_async_commit();
     };
 
-    constexpr int BUF_MASK = (Traits::STAGES > 1) ? (Traits::STAGES - 1) : 0;
+    // ---- Multi-stage cp.async pipeline ----
-
+    // Prologue loads STAGES tiles; each loop iteration waits only for the
-    // Prologue
+    // oldest outstanding group (wait_group<STAGES-1>) so the STAGES-1 newer
-    if (ti_begin < ti_end) {
+    // tile loads stay in flight and overlap with the current tile's compute.
-        load_tile(ti_begin, 0);
+    constexpr int STAGES = Traits::STAGES;
-    }
+    const int ntiles = ti_end - ti_begin;
-
-    for (int ti = ti_begin; ti < ti_end; ti++) {
-        int buf = (ti - ti_begin) & BUF_MASK;
-
-        cp_async_wait_group<0>();
-        __syncwarp();
-        if constexpr (Traits::STAGES > 1) {
-            if (ti + 1 < ti_end)
-                load_tile(ti + 1, (ti + 1 - ti_begin) & BUF_MASK);
-        }
 
+    auto process_tile = [&](int it, int buf) {
         const bf16* bK = sK + buf * Traits::BC * Traits::LD;
         const bf16* bV = sV + buf * Traits::BC * Traits::LD;
-        int kv0 = ti * Traits::BC;
+        int kv0 = (ti_begin + it) * Traits::BC;
 
         float Sacc[Traits::NC8][4];
         mma_compute_scores<Traits>(Qa, bK, lane, Sacc);
@@ -115,12 +106,29 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
                                            Sacc, Oacc, m0, m1, l0, l1, lane);
 
         mma_pv_accumulate<Traits>(Sacc, bV, lane, Oacc);
-        __syncwarp();
+    };
 
-        if constexpr (Traits::STAGES == 1) {
+    if (ntiles >= STAGES) {
-            if (ti + 1 < ti_end)
+        #pragma unroll
-                load_tile(ti + 1, 0);
+        for (int i = 0; i < STAGES; i++)
+            load_tile(ti_begin + i, i);
+
+        for (int it = 0; it < ntiles; it++) {
+            cp_async_wait_group<STAGES - 1>();
+            __syncwarp();
+            process_tile(it, it & (STAGES - 1));
+            __syncwarp();
+            if (it + STAGES < ntiles)
+                load_tile(ti_begin + it + STAGES, (it + STAGES) & (STAGES - 1));
         }
+    } else {
+        // Fewer tiles than stages: load all, wait for all, process.
+        for (int i = 0; i < ntiles; i++)
+            load_tile(ti_begin + i, i);
+        cp_async_wait_group<0>();
+        __syncwarp();
+        for (int it = 0; it < ntiles; it++)
+            process_tile(it, it);
     }
 
     // ---- write UN-normalised partials for this split ----

csrc/kernels/attn_entry_utils.cuh

@@ -21,11 +21,15 @@ using bf16 = __nv_bfloat16;
                          " (supported: 32, 64, 128, 256)"); \
     }
 
+// The split kernel unconditionally writes every (batch, q_head, split) slot it
+// owns — including empty split ranges, which store m = -FLT_MAX so the combine
+// skips them. Allocators are therefore left uninitialized (torch::empty); the
+// per-call memset (torch::zeros / torch::full) was pure overhead.
 template<typename P>
 inline void alloc_split_partials(P& p) {
     auto fopt = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);
-    auto o_part = torch::zeros(at::IntArrayRef{p.batch, p.q_head, MAX_SPLITS, p.head_dim}, fopt);
+    auto o_part = torch::empty(at::IntArrayRef{p.batch, p.q_head, MAX_SPLITS, p.head_dim}, fopt);
-    auto ml_part = torch::full(at::IntArrayRef{p.batch, p.q_head, MAX_SPLITS, 2}, -FLT_MAX, fopt);
+    auto ml_part = torch::empty(at::IntArrayRef{p.batch, p.q_head, MAX_SPLITS, 2}, fopt);
     p.o_part = (float*)o_part.data_ptr();
     p.ml_part = (float*)ml_part.data_ptr();
 }

csrc/kernels/attn_paged_decode_split_kv_mma.cuh

@@ -91,25 +91,17 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
         cp_async_commit();
     };
 
-    constexpr int BUF_MASK = (Traits::STAGES > 1) ? (Traits::STAGES - 1) : 0;
+    // ---- Multi-stage cp.async pipeline ----
-
+    // Prologue loads STAGES tiles; each loop iteration waits only for the
-    if (ti_begin < ti_end) {
+    // oldest outstanding group (wait_group<STAGES-1>) so the STAGES-1 newer
-        load_tile(ti_begin, 0);
+    // tile loads stay in flight and overlap with the current tile's compute.
-    }
+    constexpr int STAGES = Traits::STAGES;
-
+    const int ntiles = ti_end - ti_begin;
-    for (int ti = ti_begin; ti < ti_end; ti++) {
-        int buf = (ti - ti_begin) & BUF_MASK;
-
-        cp_async_wait_group<0>();
-        __syncwarp();
-        if constexpr (Traits::STAGES > 1) {
-            if (ti + 1 < ti_end)
-                load_tile(ti + 1, (ti + 1 - ti_begin) & BUF_MASK);
-        }
 
+    auto process_tile = [&](int it, int buf) {
         const bf16* bK = sK + buf * Traits::BC * Traits::LD;
         const bf16* bV = sV + buf * Traits::BC * Traits::LD;
-        int kv0 = ti * Traits::BC;
+        int kv0 = (ti_begin + it) * Traits::BC;
 
         float Sacc[Traits::NC8][4];
         mma_compute_scores<Traits>(Qa, bK, lane, Sacc);
@@ -128,12 +120,29 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
                                            Sacc, Oacc, m0, m1, l0, l1, lane);
 
         mma_pv_accumulate<Traits>(Sacc, bV, lane, Oacc);
-        __syncwarp();
+    };
 
-        if constexpr (Traits::STAGES == 1) {
+    if (ntiles >= STAGES) {
-            if (ti + 1 < ti_end)
+        #pragma unroll
-                load_tile(ti + 1, 0);
+        for (int i = 0; i < STAGES; i++)
+            load_tile(ti_begin + i, i);
+
+        for (int it = 0; it < ntiles; it++) {
+            cp_async_wait_group<STAGES - 1>();
+            __syncwarp();
+            process_tile(it, it & (STAGES - 1));
+            __syncwarp();
+            if (it + STAGES < ntiles)
+                load_tile(ti_begin + it + STAGES, (it + STAGES) & (STAGES - 1));
         }
+    } else {
+        // Fewer tiles than stages: load all, wait for all, process.
+        for (int i = 0; i < ntiles; i++)
+            load_tile(ti_begin + i, i);
+        cp_async_wait_group<0>();
+        __syncwarp();
+        for (int it = 0; it < ntiles; it++)
+            process_tile(it, it);
     }
 
     auto split_slot = [&](int h) -> size_t {

docs/guides/params.md

@@ -28,18 +28,48 @@
 | `--max_lr` | Maximum learning rate (cosine decay after warmup) | 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 `muon_adamw` optimizer sends matrix parameters through **Muon** and
+non-matrix parameters through **AdamW** (`fused=True`).
 
 | Parameter | Description | Default |
 |-----------|-------------|---------|
+| `--optimizer` | Built-in optimizer (`muon_adamw`, `nora_nadamw`, `mano_adamw`) | `muon_adamw` |
 | `--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` |
 
+`nora_nadamw` routes 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 |
+|-----------|-------------|---------|
+| `--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 |
+
+`mano_adamw` routes internal `Linear.weight` matrices to **Mano** (manifold
+normalized optimizer) and the remaining parameters to **NAdamW**. Mano projects
+the momentum onto the tangent space of the Oblique manifold and normalizes it,
+alternating the projection axis (row/column) each step — replacing Muon's
+Newton-Schulz iteration with a cheaper normalization.
+
+| Parameter | Description | Default |
+|-----------|-------------|---------|
+| `--mano_momentum` | Mano momentum factor | 0.95 |
+| `--mano_nesterov` | Enable Nesterov momentum for Mano | True |
+
+Optimizer identity and hyperparameters are saved in checkpoint metadata. Optimizer
+states are intentionally not interchangeable: resume older MuonAdamW checkpoints
+with `--optimizer=muon_adamw`.
+
 ### Data Loading
 
 | Parameter | Description | Default |

scripts/tools/train.py

@@ -1,115 +1,75 @@
 import os
+from collections import OrderedDict
 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):
+class GroupedOption(click.Option):
-    """Combined Muon (matrix) + AdamW (non-matrix) optimizer."""
+    """A ``click.Option`` that carries a ``group`` label for help output."""
 
-    def __init__(
+    def __init__(self, *args, group: str = "Options", **kwargs):
-        self,
+        super().__init__(*args, **kwargs)
-        model: nn.Module,
+        self.group = group
-        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] = []
+class GroupedCommand(click.Command):
-        for name, param in model.named_parameters():
+    """A ``click.Command`` that renders options grouped by their ``group``."""
-            if not param.requires_grad:
+
+    def format_options(self, ctx, formatter):
+        groups: OrderedDict[str, list] = OrderedDict()
+        for param in self.get_params(ctx):
+            record = param.get_help_record(ctx)
+            if record is None:
                 continue
-            if (
+            group = getattr(param, "group", "Options")
-                param.dim() >= 2
+            groups.setdefault(group, []).append(record)
-                and "norm" not in name
+        for group_name, records in groups.items():
-                and "bias" not in name
+            with formatter.section(group_name):
-                and "embed" not in name
+                formatter.write_dl(records)
-                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:
+def opt(*param_decls, group: str, **kwargs):
-    """Load YAML config, then override with explicit CLI kwargs (None excluded)."""
+    """Shorthand for ``click.option`` that tags the option with a group."""
+    kwargs.setdefault("cls", GroupedOption)
+    kwargs["group"] = group
+    return click.option(*param_decls, **kwargs)
+
+
+def _merge_yaml_into_kwargs(
+    config_path: str,
+    passed_kwargs: dict,
+    explicit_keys: set[str] | None = None,
+) -> dict:
+    """Merge Click defaults, YAML values, then explicit CLI values."""
     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,174 +77,427 @@ 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"]
 
 
 @click.command(
     name="train",
+    cls=GroupedCommand,
     help="Start model training (pretrain / SFT / DPO / GRPO).",
     context_settings={"show_default": True},
 )
-@click.option(
+@opt(
     "--config",
     "-c",
     "config_path",
     type=click.Path(exists=True),
+    group="Paths & Setup",
     help="YAML config file. CLI flags override YAML values.",
 )
-@click.option(
+@opt(
     "--train_type",
     type=click.Choice(_TRAIN_TYPE),
     required=False,
+    group="Paths & Setup",
     help="Training type.",
 )
-@click.option(
+@opt(
     "--data_root_path",
     type=click.Path(exists=True),
+    group="Paths & Setup",
     help="Root directory of the dataset.",
 )
-@click.option(
+@opt(
     "--param_path",
     type=click.Path(exists=True),
+    group="Paths & Setup",
     help="Path to model parameters or resume checkpoint.",
 )
-@click.option("--resume", is_flag=True, default=False, help="Resume from checkpoint.")
+@opt(
-@click.option("--n_epoch", type=int, default=1, help="Number of epochs.")
+    "--resume",
-@click.option("--batch_per_device", type=int, default=1, help="Batch size per GPU.")
+    is_flag=True,
-@click.option(
+    default=False,
-    "--grad_accum_steps", type=int, default=1, help="Gradient accumulation steps."
+    group="Paths & Setup",
+    help="Resume from checkpoint.",
 )
-@click.option(
+@opt("--n_epoch", type=int, default=1, group="Training", help="Number of epochs.")
+@opt(
+    "--batch_per_device",
+    type=int,
+    default=1,
+    group="Training",
+    help="Batch size per GPU.",
+)
+@opt(
+    "--grad_accum_steps",
+    type=int,
+    default=1,
+    group="Training",
+    help="Gradient accumulation steps.",
+)
+@opt(
     "--warmup_ratio",
     type=float,
     default=0.05,
+    group="LR Schedule",
     help="Fraction of total steps for LR warmup.",
 )
-@click.option("--max_lr", type=float, default=3e-4, help="Max learning rate.")
+@opt(
-@click.option(
+    "--max_lr",
-    "--max_grad_norm", type=float, default=1.0, help="Max gradient norm for clipping."
+    type=float,
+    default=3e-4,
+    group="Optimizer",
+    help="Max learning rate.",
 )
-@click.option("--weight_decay", type=float, default=0.1, help="Weight decay.")
+@opt(
-@click.option("--muon_momentum", type=float, default=0.95, help="Muon momentum factor.")
+    "--optimizer",
-@click.option("--muon_nesterov/--no-muon_nesterov", default=True, help="Muon Nesterov.")
+    type=click.Choice(_OPTIMIZERS),
-@click.option("--muon_ns_steps", type=int, default=5, help="Muon Newton-Schulz steps.")
+    default="muon_adamw",
-@click.option(
+    group="Optimizer",
+    help="Built-in optimizer.",
+)
+@opt(
+    "--max_grad_norm",
+    type=float,
+    default=1.0,
+    group="Training",
+    help="Max gradient norm for clipping.",
+)
+@opt(
+    "--weight_decay",
+    type=float,
+    default=0.1,
+    group="Optimizer",
+    help="Weight decay for eligible optimizer parameters.",
+)
+@opt(
+    "--nora_lr", type=float, default=5e-3, group="Optimizer", help="Nora learning rate."
+)
+@opt(
+    "--nora_beta", type=float, default=0.95, group="Optimizer", help="Nora EMA factor."
+)
+@opt(
+    "--nora_momentum",
+    type=float,
+    default=0.95,
+    group="Optimizer",
+    help="Nora update momentum.",
+)
+@opt(
+    "--nora_weight_decay",
+    type=float,
+    default=0.0,
+    group="Optimizer",
+    help="Nora weight decay.",
+)
+@opt(
+    "--muon_momentum",
+    type=float,
+    default=0.95,
+    group="Optimizer",
+    help="Muon momentum factor.",
+)
+@opt(
+    "--muon_nesterov/--no-muon_nesterov",
+    default=True,
+    group="Optimizer",
+    help="Muon Nesterov.",
+)
+@opt(
+    "--muon_ns_steps",
+    type=int,
+    default=5,
+    group="Optimizer",
+    help="Muon Newton-Schulz steps.",
+)
+@opt(
     "--muon_adjust_lr",
     type=click.Choice(["original", "match_rms_adamw"]),
     default="match_rms_adamw",
+    group="Optimizer",
     help="Muon LR adjustment strategy.",
 )
-@click.option("--random_seed", type=int, default=3407, help="Random seed.")
+@opt(
-@click.option("--num_workers", type=int, default=4, help="DataLoader workers.")
+    "--mano_momentum",
-@click.option("--pin_memory/--no-pin_memory", default=True, help="Pin memory.")
+    type=float,
-@click.option(
+    default=0.95,
-    "--window_size", type=int, default=None, help="Max input sequence length."
+    group="Optimizer",
+    help="Mano momentum factor.",
 )
-@click.option("--stride", type=int, default=None, help="Step size for sliding window.")
+@opt(
-@click.option("--dpo_beta", type=float, default=0.1, help="DPO beta.")
+    "--mano_nesterov/--no-mano_nesterov",
-@click.option("--group_size", type=int, default=4, help="GRPO group size.")
+    default=True,
-@click.option("--grpo_clip_eps", type=float, default=0.2, help="GRPO clip epsilon.")
+    group="Optimizer",
-@click.option(
+    help="Mano Nesterov momentum.",
-    "--grpo_kl_coef", type=float, default=0.01, help="GRPO KL penalty coefficient."
 )
-@click.option("--label_smoothing", type=float, default=0.0, help="Label smoothing.")
+@opt(
-@click.option(
+    "--random_seed",
-    "--rollout_interval", type=int, default=512, help="Steps between rollouts."
+    type=int,
+    default=3407,
+    group="Data Loading",
+    help="Random seed.",
 )
-@click.option(
+@opt(
-    "--rollout_temperature", type=float, default=0.7, help="Rollout temperature."
+    "--num_workers",
+    type=int,
+    default=4,
+    group="Data Loading",
+    help="DataLoader workers.",
 )
-@click.option("--rollout_top_k", type=int, default=0, help="Rollout top-k (0=disable).")
+@opt(
-@click.option("--rollout_top_p", type=float, default=0.9, help="Rollout top-p.")
+    "--pin_memory/--no-pin_memory",
-@click.option(
+    default=True,
+    group="Data Loading",
+    help="Pin memory.",
+)
+@opt(
+    "--window_size",
+    type=int,
+    default=None,
+    group="Data Loading",
+    help="Max input sequence length.",
+)
+@opt(
+    "--stride",
+    type=int,
+    default=None,
+    group="Data Loading",
+    help="Step size for sliding window.",
+)
+@opt("--dpo_beta", type=float, default=0.1, group="Algorithm", help="DPO beta.")
+@opt("--group_size", type=int, default=4, group="Algorithm", help="GRPO group size.")
+@opt(
+    "--grpo_clip_eps",
+    type=float,
+    default=0.2,
+    group="Algorithm",
+    help="GRPO clip epsilon.",
+)
+@opt(
+    "--grpo_kl_coef",
+    type=float,
+    default=0.01,
+    group="Algorithm",
+    help="GRPO KL penalty coefficient.",
+)
+@opt(
+    "--label_smoothing",
+    type=float,
+    default=0.0,
+    group="Data Loading",
+    help="Label smoothing.",
+)
+@opt(
+    "--rollout_interval",
+    type=int,
+    default=512,
+    group="Algorithm",
+    help="Steps between rollouts.",
+)
+@opt(
+    "--rollout_temperature",
+    type=float,
+    default=0.7,
+    group="Algorithm",
+    help="Rollout temperature.",
+)
+@opt(
+    "--rollout_top_k",
+    type=int,
+    default=0,
+    group="Algorithm",
+    help="Rollout top-k (0=disable).",
+)
+@opt(
+    "--rollout_top_p",
+    type=float,
+    default=0.9,
+    group="Algorithm",
+    help="Rollout top-p.",
+)
+@opt(
     "--rollout_max_tokens",
     type=int,
     default=1024,
+    group="Algorithm",
     help="Max tokens per rollout response.",
 )
-@click.option(
+@opt(
     "--gradient_checkpointing/--no-gradient_checkpointing",
     default=False,
+    group="Misc",
     help="Enable activation checkpointing.",
 )
-@click.option(
+@opt(
     "--compile",
     "compile_mode",
     type=click.Choice(["default", "reduce-overhead", "max-autotune"]),
     default=None,
+    group="Misc",
     help="torch.compile mode. Omit to disable.",
 )
-@click.option(
+@opt(
-    "--ckpt_interval", type=int, default=5000, help="Steps between checkpoints."
+    "--ckpt_interval",
+    type=int,
+    default=5000,
+    group="Checkpoint",
+    help="Steps between checkpoints.",
 )
-@click.option(
+@opt(
-    "--ckpt_dir", type=click.Path(), default="checkpoint", help="Checkpoint directory."
+    "--ckpt_dir",
+    type=click.Path(),
+    default="checkpoint",
+    group="Checkpoint",
+    help="Checkpoint directory.",
 )
-@click.option("--val_split", type=float, default=None, help="Validation split ratio.")
+@opt(
-@click.option(
+    "--val_split",
-    "--val_step", type=int, default=1000, help="Steps between validation runs."
+    type=float,
+    default=None,
+    group="Validation",
+    help="Validation split ratio.",
 )
-@click.option(
+@opt(
+    "--val_step",
+    type=int,
+    default=1000,
+    group="Validation",
+    help="Steps between validation runs.",
+)
+@opt(
     "--metrics",
     multiple=True,
-    default=("loss", "lr", "grad_norm"),
+    default=("loss", "lr", "grad_norm", "grad_snr"),
+    group="Validation",
     help="Metrics to log (repeatable).",
 )
-@click.option("--start_epoch", type=int, default=0, help="Start epoch.")
+@opt("--start_epoch", type=int, default=0, group="Checkpoint", help="Start epoch.")
-@click.option("--start_samples", type=int, default=0, help="Start samples (per rank).")
+@opt(
-@click.option(
+    "--start_samples",
-    "--master_addr", type=str, default="localhost", help="Master node address."
+    type=int,
+    default=0,
+    group="Checkpoint",
+    help="Start samples (per rank).",
 )
-@click.option("--master_port", type=str, default="29500", help="Master node port.")
+@opt(
-@click.option(
+    "--master_addr",
+    type=str,
+    default="localhost",
+    group="Distributed",
+    help="Master node address.",
+)
+@opt(
+    "--master_port",
+    type=str,
+    default="29500",
+    group="Distributed",
+    help="Master node port.",
+)
+@opt(
     "--backend",
     type=click.Choice(_BACKENDS),
     default="nccl",
+    group="Distributed",
     help="Distributed backend.",
 )
-@click.option("--nprocs", type=int, default=1, help="Number of GPUs.")
+@opt("--nprocs", type=int, default=1, group="Distributed", help="Number of GPUs.")
-@click.option(
+@opt(
     "--parallel_mode",
     type=click.Choice(_PARALLEL),
     default="fsdp",
+    group="Distributed",
     help="Parallel strategy.",
 )
-@click.option("--device_type", type=str, default="cuda", help="Device type.")
+@opt(
-@click.option(
+    "--device_type",
+    type=str,
+    default="cuda",
+    group="Distributed",
+    help="Device type.",
+)
+@opt(
     "--start_method",
     type=click.Choice(_START_METHODS),
     default="spawn",
+    group="Distributed",
     help="Multiprocessing start method.",
 )
-@click.option("--neftune_alpha", type=float, default=0.0, help="NEFTune noise alpha.")
+@opt(
-@click.option(
+    "--neftune_alpha",
+    type=float,
+    default=0.0,
+    group="Algorithm",
+    help="NEFTune noise alpha.",
+)
+@opt(
     "--schedule_type",
     type=click.Choice(_SCHEDULES),
     default="cosine",
+    group="LR Schedule",
     help="LR scheduler.",
 )
-@click.option(
+@opt(
-    "--min_rate", type=float, default=None, help="Minimum LR as fraction of base LR."
+    "--min_rate",
+    type=float,
+    default=None,
+    group="LR Schedule",
+    help="Minimum LR as fraction of base LR.",
 )
-@click.option("--cycle_length", type=int, default=None, help="SGDR first cycle length.")
+@opt(
-@click.option("--t_mult", type=int, default=2, help="SGDR cycle length multiplier.")
+    "--cycle_length",
-@click.option(
+    type=int,
-    "--stable_steps", type=int, default=None, help="WSD stable plateau steps."
+    default=None,
+    group="LR Schedule",
+    help="SGDR first cycle length.",
 )
-@click.option("--decay_steps", type=int, default=None, help="WSD decay steps.")
+@opt(
-@click.option("--tp_size", type=int, default=None, help="Tensor parallelism (future).")
+    "--t_mult",
-@click.option(
+    type=int,
+    default=2,
+    group="LR Schedule",
+    help="SGDR cycle length multiplier.",
+)
+@opt(
+    "--stable_steps",
+    type=int,
+    default=None,
+    group="LR Schedule",
+    help="WSD stable plateau steps.",
+)
+@opt(
+    "--decay_steps",
+    type=int,
+    default=None,
+    group="LR Schedule",
+    help="WSD decay steps.",
+)
+@opt(
+    "--tp_size",
+    type=int,
+    default=None,
+    group="Distributed",
+    help="Tensor parallelism (future).",
+)
+@opt(
     "--dry-run",
     is_flag=True,
     default=False,
+    group="Misc",
     help="Validate config and print plan, do not train.",
 )
 @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 +508,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 +530,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", "muon_adamw"))),
         ("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 +550,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 = "muon_adamw", **kwargs
+) -> optim.Optimizer:
+    return OptimizerFactory.create(optimizer_name, model, **kwargs)
 
 
 def create_scheduler(
@@ -459,15 +675,61 @@ def train(
         tokenizer_path=param_path,
     )
 
+    optimizer_name = kwargs.pop("optimizer", "muon_adamw")
+    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"),
+        "mano_momentum": kwargs.pop("mano_momentum", 0.95),
+        "mano_nesterov": kwargs.pop("mano_nesterov", True),
+    }
     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}
+        )
+    elif optimizer_name == "mano_adamw":
+        optimizer_hyperparameters = {
+            key: optimizer_kwargs[key]
+            for key in ("lr", "weight_decay", "mano_momentum", "mano_nesterov")
+        }
+        optimizer_hyperparameters.update(
+            {"adamw_betas": [0.9, 0.95], "adamw_eps": 1e-8}
+        )
+    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 +778,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_mano_adamw.py

@@ -0,0 +1,119 @@
+import math
+from copy import deepcopy
+
+import pytest
+import torch
+
+from astrai.optim import Mano, ManoAdamW, OptimizerFactory
+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_mano_one_step_projects_to_tangent_space():
+    original = torch.tensor([[3.0, 4.0], [0.0, 2.0]])
+    param = torch.nn.Parameter(original.clone())
+    grad = torch.tensor([[4.0, -3.0], [1.0, 1.0]])
+    param.grad = grad.clone()
+
+    optimizer = Mano(
+        [param], lr=0.1, momentum=0.0, nesterov=False, eps=1e-8, weight_decay=0.0
+    )
+    optimizer.step()
+
+    dim = 0
+    tangent = grad - (torch.sum(grad * original, dim=dim, keepdim=True) * original)
+    direction = tangent / (torch.norm(tangent, p=2, dim=dim, keepdim=True) + 1e-8)
+    adjusted_lr = 0.1 * 0.2 * math.sqrt(direction.shape[dim])
+    expected = original - adjusted_lr * direction
+    torch.testing.assert_close(param, expected)
+
+
+def test_mano_alternates_projection_axis():
+    param = torch.nn.Parameter(torch.eye(4) * 3.0)
+    param.grad = torch.ones(4, 4)
+
+    optimizer = Mano([param], lr=0.1, momentum=0.0, nesterov=False)
+    optimizer.step()
+    dim_step0 = 0
+
+    param.grad = torch.ones(4, 4)
+    optimizer.step()
+    dim_step1 = 1
+
+    assert dim_step0 != dim_step1
+
+
+def test_mano_rejects_non_2d_parameters():
+    param = torch.nn.Parameter(torch.randn(3, 4, 5))
+    with pytest.raises(ValueError, match="2D"):
+        Mano([param])
+
+
+def test_factory_registers_mano():
+    assert "mano_adamw" in OptimizerFactory.list_registered()
+    from astrai.model import AutoRegressiveLM
+
+    model = AutoRegressiveLM(make_tiny_config())
+    optimizer = OptimizerFactory.create("mano_adamw", model, lr=3e-4)
+    assert isinstance(optimizer, ManoAdamW)
+
+
+def test_mano_adamw_runs_closure_once():
+    from astrai.model import AutoRegressiveLM
+
+    model = AutoRegressiveLM(make_tiny_config())
+    optimizer = ManoAdamW(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
+
+
+def test_mano_adamw_resume_matches_uninterrupted():
+    from astrai.model import AutoRegressiveLM
+    from astrai.trainer.schedule import SchedulerFactory
+
+    torch.manual_seed(7)
+    model_a = AutoRegressiveLM(make_tiny_config())
+    optimizer_a = ManoAdamW(model_a, lr=3e-4)
+    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 = ManoAdamW(model_b, lr=3e-4)
+    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())

tests/optim/test_nora_nadamw.py

@@ -0,0 +1,257 @@
+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() == [
+        "mano_adamw",
+        "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)