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

@@ -1,17 +1,16 @@
 """Rotary embedding with auto-dispatch to CUDA kernel.
 
-Single entry point ``apply_rotary_emb(x, cos, sin)`` — uses the fused
+Single entry point ``apply_rotary_emb(x, freqs_cis)`` — uses the fused
 CUDA kernel when available, falls back to torch complex multiply otherwise.
 
 Layout: x is [batch, seq_len, n_heads, head_dim] (bf16).
-cos/sin are [batch, seq_len, head_dim/2] (f32).
+freqs_cis is [batch, seq_len, dim/2, 2] (f32) — [cos, sin] pairs.
 """
 
 import torch
 from torch import Tensor
 
 from astrai.extension.loader import is_available
-from astrai.extension.rotary_ops import rotary_emb as _cuda_rotary
 
 _cache = {"available": None}
 
@@ -22,32 +21,34 @@ def _cuda_available() -> bool:
     return _cache["available"]
 
 
-def _torch_apply(x: Tensor, cos: Tensor, sin: Tensor) -> Tensor:
+def _torch_apply(x: Tensor, freqs_cis: Tensor) -> Tensor:
+    cos, sin = freqs_cis[..., 0], freqs_cis[..., 1]
     dtype = x.dtype
     x_ = x.float().reshape(*x.shape[:-1], -1, 2)
     x_complex = torch.view_as_complex(x_)
-    freqs_cis = torch.complex(cos, sin).unsqueeze(2)
+    freqs_cis_complex = torch.complex(cos, sin).unsqueeze(2)
-    x_rotated = x_complex * freqs_cis
+    x_rotated = x_complex * freqs_cis_complex
     x_out = torch.view_as_real(x_rotated).flatten(-2)
     return x_out.to(dtype)
 
 
-def apply_rotary_emb(x: Tensor, rotary_emb: tuple[Tensor, Tensor]) -> Tensor:
+def apply_rotary_emb(x: Tensor, freqs_cis: Tensor) -> Tensor:
     """Apply rotary embedding to x.
 
     Args:
         x: [batch, seq_len, n_heads, head_dim] (bf16)
-        rotary_emb: (cos, sin) tuple, each [batch, seq_len, head_dim/2] (f32)
+        freqs_cis: [batch, seq_len, dim/2, 2] (f32) — [cos, sin] pairs
 
     Returns:
         [batch, seq_len, n_heads, head_dim] (bf16)
     """
-    cos, sin = rotary_emb
     if (
         _cuda_available()
         and not torch.is_grad_enabled()
         and x.is_cuda
         and x.dtype == torch.bfloat16
     ):
-        return _cuda_rotary(x, cos, sin)
+        from astrai.extension.rotary_ops import rotary_emb as _cuda_rotary
-    return _torch_apply(x, cos, sin)
+
+        return _cuda_rotary(x, freqs_cis)
+    return _torch_apply(x, freqs_cis)

astrai/extension/rotary_ops.py

@@ -2,10 +2,10 @@
 
 Calls the compiled CUDA kernel directly. If the kernel is not available,
 raises ``RuntimeError``. Fallback to torch complex multiply is the
-responsibility of ``astrai.model.components.rope.apply_rotary_emb``.
+responsibility of ``astrai.extension.rotary_backend.apply_rotary_emb``.
 
-Layout convention: x is ``[batch, seq_len, n_heads, head_dim]`` (blhd, bf16).
+Layout: x is [batch, seq_len, n_heads, head_dim] (bf16, contiguous).
-cos/sin are ``[batch, seq_len, head_dim/2]`` (f32).
+freqs_cis is [batch, seq_len, head_dim/2, 2] (f32, contiguous) — [cos, sin] pairs.
 """
 
 import torch
@@ -21,21 +21,12 @@ def _check_available():
         )
 
 
-def rotary_emb(
+def rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
-    x: torch.Tensor,
-    cos: torch.Tensor,
-    sin: torch.Tensor,
-) -> torch.Tensor:
     """Fused rotary embedding kernel.
 
-    Applies rotation: for each pair (x_even, x_odd):
-        out_even = x_even * cos - x_odd * sin
-        out_odd  = x_even * sin + x_odd * cos
-
     Args:
         x: [batch, seq_len, n_heads, head_dim] (bf16, contiguous)
-        cos: [batch, seq_len, head_dim/2] (f32)
+        freqs_cis: [batch, seq_len, head_dim/2, 2] (f32, contiguous) — [cos, sin] pairs
-        sin: [batch, seq_len, head_dim/2] (f32)
 
     Returns:
         [batch, seq_len, n_heads, head_dim] (bf16)
@@ -43,4 +34,6 @@ def rotary_emb(
     _check_available()
     if not x.is_contiguous():
         x = x.contiguous()
-    return _modules["rotary_emb"].rotary_emb(x, cos, sin)
+    if not freqs_cis.is_contiguous():
+        freqs_cis = freqs_cis.contiguous()
+    return _modules["rotary_emb"].rotary_emb(x, freqs_cis)

astrai/model/components/rope.py

@@ -1,4 +1,4 @@
-from typing import Dict, Optional, Tuple
+from typing import Dict, Optional
 
 import torch
 import torch.nn as nn
@@ -10,16 +10,18 @@ def get_rotary_emb(
     max_len: int,
     base: float = 10000,
     device: Optional[torch.device] = None,
-) -> Tuple[Tensor, Tensor]:
+) -> Tensor:
     """Precompute cos/sin tables for rotary embedding.
 
     Returns:
-        (cos, sin) each of shape [max_len, dim/2] (f32)
+        [max_len, dim/2, 2] (f32) — [cos, sin] pairs.
     """
     theta = base ** (-torch.arange(0, dim, 2, dtype=torch.float64, device=device) / dim)
     t = torch.arange(0, max_len, dtype=torch.float64, device=device)
     freqs = torch.outer(t, theta).float()
-    return torch.cos(freqs), torch.sin(freqs)
+    cos = torch.cos(freqs)
+    sin = torch.sin(freqs)
+    return torch.stack([cos, sin], dim=-1)
 
 
 def ntk_base(base: float, dim: int, factor: float) -> float:
@@ -49,13 +51,10 @@ class RotaryEmbedding(nn.Module):
         self._set_rotary_buffer(self.max_len)
 
     def _set_rotary_buffer(self, max_len: int):
-        cos, sin = get_rotary_emb(self.dim, max_len, self.base)
+        freqs_cis = get_rotary_emb(self.dim, max_len, self.base)
-        self.register_buffer("cos_table", cos, persistent=False)
+        self.register_buffer("freqs_cis", freqs_cis, persistent=False)
-        self.register_buffer("sin_table", sin, persistent=False)
 
-    def forward(
+    def forward(self, x: Tensor, position_ids: Optional[Tensor] = None) -> Tensor:
-        self, x: Tensor, position_ids: Optional[Tensor] = None
-    ) -> Tuple[Tensor, Tensor]:
         """Lookup cos/sin for the given positions.
 
         Args:
@@ -63,7 +62,7 @@ class RotaryEmbedding(nn.Module):
             position_ids: [batch, seq_len] optional position indices.
 
         Returns:
-            (cos, sin) each of shape [batch, seq_len, dim/2] (f32)
+            [batch, seq_len, dim/2, 2] (f32) — [cos, sin] pairs.
         """
         if position_ids is None:
             position_ids = (
@@ -71,6 +70,4 @@ class RotaryEmbedding(nn.Module):
                 .unsqueeze(0)
                 .expand(x.size(0), -1)
             )
-        cos = self.cos_table[position_ids].float()
+        return self.freqs_cis[position_ids].float()
-        sin = self.sin_table[position_ids].float()
-        return cos, sin

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/build.py

@@ -53,7 +53,7 @@ NVCC_FLAGS = [
     "--use_fast_math",
     "--ptxas-options=-O3,-v",
     "--extra-device-vectorization",
-    "--threads=8",
+    "--threads=16",
 ]
 
 

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

@@ -1,4 +1,5 @@
 #pragma once
+#include <float.h>
 #include <torch/extension.h>
 #include <c10/cuda/CUDAGuard.h>
 #include "attn_common.h"
@@ -20,6 +21,10 @@ 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);

csrc/kernels/attn_paged_decode_split_kv.cuh

@@ -67,14 +67,17 @@ __global__ void paged_attn_decode_split_kv_kernel(PagedAttentionParams<bf16> p)
             partial = warp_reduce_sum(partial) * p.scale;
 
             int kv_idx = chunk_start + s;
+            bool masked = false;
             if constexpr (HasMask) {
                 if (!p.mask[mask_base + kv_idx])
-                    partial = -FLT_MAX;
+                    masked = true;
             }
             if constexpr (IsCausal) {
                 if (kv_idx > p.causal_offset)
-                    partial = -FLT_MAX;
+                    masked = true;
             }
+            if (masked)
+                partial = -FLT_MAX;
 
             float new_m = fmaxf(m, partial);
             float alpha = __expf(m - new_m);
@@ -85,7 +88,11 @@ __global__ void paged_attn_decode_split_kv_kernel(PagedAttentionParams<bf16> p)
             int logical_page = pos / p.page_size;
             int page_offset = pos % p.page_size;
             int phys_page = p.page_table[batch * p.max_pages + logical_page];
-            if (phys_page >= 0) {
+            if (masked) {
+                #pragma unroll
+                for (int i = 0; i < hd_per_thread; i++)
+                    acc_reg[i] = fmaf(acc_reg[i], alpha, 0.0f);
+            } else if (phys_page >= 0) {
                 int64_t v_base = (int64_t)phys_page * p.page_size * p.kv_head * p.head_dim
                                + (int64_t)page_offset * p.kv_head * p.head_dim
                                + (int64_t)kv_head * p.head_dim;

csrc/kernels/attn_paged_decode_split_kv_mma.cuh

@@ -31,6 +31,13 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
     __shared__ __align__(16) bf16 sK[Traits::STAGES * Traits::BC * Traits::LD];
     __shared__ __align__(16) bf16 sV[Traits::STAGES * Traits::BC * Traits::LD];
 
+    #pragma unroll
+    for (int i = lane; i < Traits::STAGES * Traits::BC * Traits::LD; i += 32) {
+        sK[i] = __float2bfloat16(0.0f);
+        sV[i] = __float2bfloat16(0.0f);
+    }
+    __syncwarp();
+
     const int q_base = batch * p.q_stride_b + q_head0 * p.q_stride_h;
     const int qra = gid;
     const int qrb = gid + 8;
@@ -68,6 +75,9 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
             int r = i / Traits::HEAD_DIM, d = i % Traits::HEAD_DIM;
             int kc = kv0 + r;
             bool valid = (kc < p.kv_len);
+            if constexpr (HasMask) {
+                valid = valid && p.mask[batch * p.mask_b_stride + kc];
+            }
             int phys_page = valid ? p.page_table[batch * p.max_pages + kc] : 0;
             valid = valid && (phys_page >= 0);
             int page_off = kc % p.page_size;
@@ -81,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);
@@ -118,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 {

csrc/kernels/rotary_emb.cu

@@ -3,8 +3,7 @@
 
 __global__ void rotary_emb_kernel(
     const __nv_bfloat16* __restrict__ x,
-    const float* __restrict__ cos,
+    const float* __restrict__ freqs_cis,
-    const float* __restrict__ sin,
     __nv_bfloat16* __restrict__ out,
     int batch,
     int seq_len,
@@ -26,14 +25,14 @@ __global__ void rotary_emb_kernel(
         int b = tmp / seq_len;
 
         int x_offset = ((b * seq_len + seq) * n_heads + head) * head_dim + (pair << 1);
-        int cs_offset = (b * seq_len + seq) * half_dim + pair;
+        int cs_offset = ((b * seq_len + seq) * half_dim + pair) * 2;
 
         __nv_bfloat162 x_pair = *reinterpret_cast<const __nv_bfloat162*>(x + x_offset);
         float x_even = __bfloat162float(__low2bfloat16(x_pair));
         float x_odd  = __bfloat162float(__high2bfloat16(x_pair));
 
-        float c = cos[cs_offset];
+        float c = freqs_cis[cs_offset];
-        float s = sin[cs_offset];
+        float s = freqs_cis[cs_offset + 1];
 
         float out_even = x_even * c - x_odd * s;
         float out_odd  = x_even * s + x_odd * c;
@@ -45,23 +44,21 @@ __global__ void rotary_emb_kernel(
 
 torch::Tensor rotary_emb(
     torch::Tensor x,
-    torch::Tensor cos,
+    torch::Tensor freqs_cis
-    torch::Tensor sin
 ) {
+    TORCH_CHECK(x.is_cuda(), "x must be on CUDA");
+    TORCH_CHECK(freqs_cis.is_cuda(), "freqs_cis must be on CUDA");
+    TORCH_CHECK(x.scalar_type() == torch::kBFloat16, "x must be bf16");
+    TORCH_CHECK(x.dim() == 4, "x must be 4D [batch, seq_len, n_heads, head_dim]");
+    TORCH_CHECK(x.is_contiguous(), "x must be contiguous");
+    TORCH_CHECK(freqs_cis.dim() == 4, "freqs_cis must be 4D [batch, seq_len, dim/2, 2]");
+    TORCH_CHECK(freqs_cis.is_contiguous(), "freqs_cis must be contiguous");
 
     int batch = x.size(0);
     int seq_len = x.size(1);
     int n_heads = x.size(2);
     int head_dim = x.size(3);
 
-    TORCH_CHECK(x.is_cuda(), "x must be on CUDA");
-    TORCH_CHECK(cos.is_cuda(), "cos must be on CUDA");
-    TORCH_CHECK(sin.is_cuda(), "sin must be on CUDA");
-    TORCH_CHECK(x.scalar_type() == torch::kBFloat16, "x must be bf16");
-    TORCH_CHECK(x.dim() == 4, "x must be 4D [batch, seq_len, n_heads, head_dim]");
-    TORCH_CHECK(x.is_contiguous(), "x must be contiguous");
-    TORCH_CHECK(cos.dim() == 3, "cos must be 3D [batch, seq_len, head_dim/2]");
-    TORCH_CHECK(sin.dim() == 3, "sin must be 3D [batch, seq_len, head_dim/2]");
     TORCH_CHECK(head_dim % 2 == 0, "head_dim must be even");
 
     auto out = torch::empty_like(x);
@@ -73,8 +70,7 @@ torch::Tensor rotary_emb(
 
     rotary_emb_kernel<<<grid, block>>>(
         reinterpret_cast<const __nv_bfloat16*>(x.data_ptr()),
-        cos.data_ptr<float>(),
+        freqs_cis.data_ptr<float>(),
-        sin.data_ptr<float>(),
         reinterpret_cast<__nv_bfloat16*>(out.data_ptr()),
         batch, seq_len, n_heads, head_dim
     );
@@ -85,8 +81,7 @@ torch::Tensor rotary_emb(
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
     m.def("rotary_emb", &rotary_emb,
         py::arg("x"),
-        py::arg("cos"),
+        py::arg("freqs_cis"),
-        py::arg("sin"),
+        "Fused rotary embedding (bf16 x, f32 freqs_cis [b,s,d/2,2], bf16 out)"
-        "Fused rotary embedding (bf16 x, f32 cos/sin, bf16 out)"
     );
 }

docs/developer/architecture.md

@@ -380,7 +380,9 @@ classDiagram
             +int max_len
             +float base
             +Optional[Dict] rope_scaling
-            +forward(x, position_ids=None) Tensor
+            +Tensor cos_table
+            +Tensor sin_table
+            +forward(x, position_ids=None) Tuple[Tensor, Tensor]
         }
 
         class Embedding {
@@ -849,6 +851,9 @@ classDiagram
             +Tensor req_pool_indices
             +Tensor seq_lens
             +Tensor out_cache_loc
+            +int max_len
+            +Optional[Tensor] page_table
+            +Optional[Tensor] decode_mask
         }
 
         class PagePool {
@@ -1401,7 +1406,7 @@ classDiagram
 | **astrai.tokenize** | AutoTokenizer, ChatTemplate | Tokenizer and chat template |
 | **astrai.trainer** | Trainer, TrainContext, TrainContextBuilder, BaseStrategy–GRPOStrategy, StrategyFactory, BaseScheduler–WSDScheduler, SchedulerFactory, TrainCallback(Protocol)–MetricCallback, CallbackFactory, RawRollout, RolloutResult, BaseRewardModel, RolloutGenerator, RolloutRunner | Training workflow |
 | **astrai.inference** | InferenceEngine, InferenceScheduler, Executor, PagePool, KVStorage, ReqToTokenPool, KVCache, Allocator, PrefixCache, Task, TaskManager, TaskStatus, StreamDecoder, GenerationRequest, GenerateResult, BaseSamplingStrategy–SamplingPipeline, FrequencyPenaltyStrategy, ProtocolHandler, ResponseBuilder, OpenAIResponseBuilder, AnthropicResponseBuilder, StopChecker, GenContext, StopInfo, ChatMessage, FunctionDef, ToolDef, ChatCompletionRequest, AnthropicMessage, MessagesRequest, BaseToolParser, ToolParserFactory, SimpleJsonToolParser | Inference service |
-| **astrai.extension** | AttentionBackend, TorchNativeBackend, CudaBackend, attn_backend, ATTN_BACKEND, attn_decode, attn_prefill, attn_paged_decode, is_available | CUDA attention kernels + backend abstraction |
+| **astrai.extension** | AttentionBackend, TorchNativeBackend, CudaBackend, attn_backend, ATTN_BACKEND, attn_decode, attn_prefill, attn_paged_decode, rotary_emb, apply_rotary_emb, rotary_backend, is_available | CUDA attention + rotary kernels, backend abstraction, auto-dispatch |
 | **astrai.parallel** | spawn_parallel_fn, setup_parallel, get_rank/get_world_size/get_current_device, only_on_rank, LaunchStrategy, TorchrunStrategy, LocalStrategy, BaseExecutor, ExecutorFactory, NoneExecutor, DDPExecutor, FSDPExecutor, GradientState, AccumOptimizer, AccumScheduler | Distributed parallel & gradient accumulation |
 | **astrai.factory** | BaseFactory | Component registration |
 | **astrai.protocols** | OptimizerProtocol, SchedulerProtocol | Structural subtyping for optimizer/scheduler wrappers |
@@ -1420,6 +1425,7 @@ classDiagram
 | **Context** | `TrainContext` | Unified training state bag |
 | **Object Pool** | `Allocator`, `PagePool` | Page-based KV cache with LRU eviction |
 | **Strategy (Attention)** | `AttentionBackend`, `TorchNativeBackend`, `CudaBackend` | Attention computation backend switching via context manager |
+| **Auto-dispatch (Rotary)** | `apply_rotary_emb`, `rotary_backend.py`, `rotary_ops.py` | Rotary embedding CUDA kernel auto-dispatch with torch fallback |
 | **Executor** | `BaseExecutor`, `NoneExecutor`, `DDPExecutor`, `FSDPExecutor` | Gradient accumulation & model distribution |
 | **Storage** | `Store`, `H5Store`, `MmapStore`, `JsonlStore` | Format-agnostic data access with multi-segment support |
 | **Producer-Consumer** | `InferenceScheduler`, `Task`, queues | Continuous batching |
@@ -1431,7 +1437,7 @@ classDiagram
 2. **Training Flow**: `Trainer` → `TrainContextBuilder` → `TrainContext`, uses `BaseStrategy` for loss, `BaseExecutor` for gradient accumulation + model distribution
 3. **Strategy Selection**: `StrategyFactory` creates strategy by `train_type`
 4. **Executor Selection**: `ExecutorFactory.create(cfg.parallel_mode, grad_accum_steps=cfg.grad_accum_steps, **cfg.executor_kwargs)` → `NoneExecutor` / `DDPExecutor` / `FSDPExecutor`
-5. **Inference Flow**: `InferenceEngine` → `InferenceScheduler` → `AutoRegressiveLM`, backed by `PagePool` + `KVCache` + `SamplingPipeline`. Attention backend selected via `attn_backend()` context manager (`TorchNativeBackend` default, `CudaBackend` for CUDA kernels).
+5. **Inference Flow**: `InferenceEngine` → `InferenceScheduler` → `AutoRegressiveLM`, backed by `PagePool` + `KVCache` + `SamplingPipeline`. Attention backend selected via `attn_backend()` context manager (`TorchNativeBackend` default, `CudaBackend` for CUDA kernels). Rotary embedding auto-dispatches to CUDA kernel when available (inference mode), else torch complex multiply (training).
 6. **Distributed**: `spawn_parallel_fn` + `setup_parallel` for multi-process DDP
 7. **Dataset Loading**: `DatasetFactory` creates datasets, `Store` (H5Store/MmapStore/JsonlStore) loads data with explicit `_length` and multi-segment `_data`
 8. **Checkpoint**: `Checkpoint` saves/loads safetensors + metadata (rank-0 only), extra state saved as `{key}.pt`
@@ -1439,4 +1445,4 @@ classDiagram
 10. **AutoModel**: `from_pretrained()` loads `config.json` + `model.safetensors`, `_disable_random_init` replaces `nn.init.*` with no-ops
 11. **Protocols**: `OptimizerProtocol` / `SchedulerProtocol` — structural subtyping for `AccumOptimizer` / `AccumScheduler` wrappers
 
-> Document Update Time: 2026-07-30
+> Document Update Time: 2026-07-31

docs/developer/cuda_kernels.md

@@ -1,6 +1,6 @@
 # CUDA Kernels
 
-AstrAI includes optional custom CUDA attention kernels for decode and prefill. These are built when `nvcc` is available and CUDA is detected, and are dispatched via the `CudaBackend` attention backend.
+AstrAI includes optional custom CUDA kernels for attention and rotary embedding. These are built when `nvcc` is available and CUDA is detected, and are dispatched via the `CudaBackend` attention backend or auto-dispatched for rotary.
 
 ## Overview
 
@@ -9,6 +9,7 @@ AstrAI includes optional custom CUDA attention kernels for decode and prefill. T
 | `attn_decode` | `attn_decode.cu` | GQA decode attention (split-KV) |
 | `attn_prefill` | `attn_prefill.cu` | GQA prefill attention (split-Q) |
 | `attn_paged_decode` | `attn_paged_decode.cu` | Paged KV cache decode attention |
+| `rotary_emb` | `rotary_emb.cu` | Fused rotary embedding (cos/sin lookup + rotation) |
 
 Additionally, optimized `.cuh` variants with tensor-core MMA (Matrix Multiply-Accumulate) exist:
 
@@ -18,6 +19,18 @@ Additionally, optimized `.cuh` variants with tensor-core MMA (Matrix Multiply-Ac
 | Split-Q MMA prefill | `attn_prefill_split_q_mma.cuh` | Split Q across warps + MMA (sm_80+) |
 | Paged split-KV MMA decode | `attn_paged_decode_split_kv_mma.cuh` | Paged cache + split-KV + MMA |
 
+### Rotary Embedding Kernel
+
+The `rotary_emb` kernel (`csrc/kernels/rotary_emb.cu`) fuses cos/sin lookup and rotation into a single kernel:
+
+- One thread per (head, dim-pair), vectorized `__nv_bfloat162` load/store
+- f32 cos/sin input, bf16 compute and output
+- 256-thread blocks, grid-stride loop
+- Auto-dispatched via `apply_rotary_emb` in `astrai/extension/rotary_backend.py` (CUDA when available + inference mode, else torch complex-multiply fallback)
+- No context-manager backend needed — rotary is backend-agnostic, both attention backends benefit
+
+Standalone benchmark vs torch complex-multiply (48 calls = 24 layers × q+k): 6-9x faster, max diff 0 (decode) to 3e-2 (large prefill, bf16).
+
 ## Build System
 
 ### Auto-detection
@@ -36,7 +49,7 @@ CSRC_KERNELS=true pip install -e . --no-build-isolation
 
 # Rebuild after editing .cu/.cuh files
 CSRC_KERNELS=true python setup.py build_ext --inplace
-# Output: astrai/extension/*.so
+# Output: astrai/extension/lib/*.so
 ```
 
 ### Architecture flags
@@ -53,7 +66,7 @@ NVCC_FLAGS = -O3 --expt-relaxed-constexpr --use_fast_math
              --ptxas-options=-O3,-v --extra-device-vectorization --threads=8
 ```
 
-The `REGISTRY` in `csrc/build.py` lists all registered kernels (currently 3). Each entry maps a kernel name to its source files and build flags.
+The `REGISTRY` in `csrc/build.py` lists all registered kernels (currently 4). Each entry maps a kernel name to its source files and build flags.
 
 ## Attention Backend
 
@@ -74,9 +87,20 @@ with attn_backend(ATTN_BACKEND.CUDA):
 
 `CudaBackend` falls back to `TorchNativeBackend` when a kernel is not available.
 
+### Rotary Backend
+
+`astrai/extension/rotary_backend.py` provides `apply_rotary_emb(x, (cos, sin))` with auto-dispatch:
+
+- **CUDA path**: calls `rotary_emb` kernel directly when available, input is bf16 on CUDA, and `torch.is_grad_enabled()` is `False` (inference)
+- **Torch fallback**: complex multiply (`torch.view_as_complex` → `torch.complex` multiply → `torch.view_as_real`), used during training (supports autograd) or when kernel unavailable
+
+No context-manager switching needed — the dispatch is automatic per call.
+
 ## Python Wrappers
 
-`astrai/extension/attention_ops.py` provides Python wrappers for each compiled kernel. Each wrapper calls its CUDA kernel directly and raises `RuntimeError` if the `.so` is not available. Fallback to torch SDPA is handled by the attention backend, not the wrapper functions.
+`astrai/extension/attention_ops.py` provides Python wrappers for each compiled attention kernel. Each wrapper calls its CUDA kernel directly and raises `RuntimeError` if the `.so` is not available. Fallback to torch SDPA is handled by the attention backend, not the wrapper functions.
+
+`astrai/extension/rotary_ops.py` provides the wrapper for the rotary embedding kernel. Fallback to torch complex multiply is handled by `rotary_backend.py`.
 
 Interface (all functions):
 ```
@@ -115,8 +139,8 @@ nvcc -I csrc -arch=sm_89 -O3 --use_fast_math \
 ## Known Optimization Targets
 
 - **Decode D=256**: spill eliminated (BC=16 + STAGES=2), but still 248 regs — further tiling could help.
-- **Prefill single-batch**: bandwidth low (52 GB/s at q=kv=2048) — likely compute-bound but near L20 bf16 ceiling (~94 TFLOP/s).
+- **Prefill single-batch**: bandwidth low (22 GB/s at q=kv=2048) — compute-bound at ~94 TFLOP/s (near L20 bf16 ceiling ~193 TFLOP/s for non-causal).
-- **Decode single-batch**: bandwidth low (309 GB/s at kv=512) — L20 HBM ~864 GB/s theoretical; small kv underutilizes SMs despite split-KV.
+- **Decode single-batch**: bandwidth low (113 GB/s at kv=512, 13% of 864 GB/s theoretical) — small kv underutilizes SMs despite split-KV; scales to 757 GB/s (88%) at B=16+.
 
 ## File Layout
 
@@ -124,10 +148,11 @@ nvcc -I csrc -arch=sm_89 -O3 --use_fast_math \
 csrc/
 ├── build.py              # Build system: REGISTRY, _arch_flags, nvcc flags
 ├── kernels/
-│   ├── attn_common.h     # Shared attention utilities
+│   ├── attn_common.h     # Shared attention params (AttentionParams, PagedAttentionParams)
 │   ├── attn_decode.cu    # Basic decode kernel (registered)
 │   ├── attn_prefill.cu   # Basic prefill kernel (registered)
 │   ├── attn_paged_decode.cu             # Paged decode kernel (registered)
+│   ├── rotary_emb.cu                    # Fused rotary embedding kernel (registered)
 │   ├── attn_decode_split_kv.cuh         # Split-KV variant
 │   ├── attn_decode_split_kv_mma.cuh     # Split-KV + MMA variant
 │   ├── attn_prefill_split_q.cuh         # Split-Q variant
@@ -145,4 +170,6 @@ csrc/
     └── attn_prefill_test.cu     # Prefill kernel test
 ```
 
-> Document Update Time: 2026-07-30
+Compiled `.so` files are placed in `astrai/extension/lib/`, separate from Python source files.
+
+> Document Update Time: 2026-07-31

docs/developer/internals.md

@@ -41,7 +41,7 @@ RoPE embeds position into Q/K vectors via complex rotation:
 
 $$ q_i = R_i W_q x_i, \quad k_j = R_j W_k x_j, \quad q_i^T k_j = x_i^T W_q^T R_{i-j} W_k x_j $$
 
-The complex rotation `freqs_cis` is pre-computed once (`cos, sin` pairs per position). `apply_rotary_emb` multiplies Q/K as complex numbers. The key property is that the dot product $q_i^T k_j$ depends only on the relative position $i - j$, not the absolute positions.
+`RotaryEmbedding` pre-computes `cos_table` and `sin_table` (f32, `[max_len, dim/2]`). `forward()` returns a `(cos, sin)` tuple indexed by `position_ids`. `apply_rotary_emb` applies the rotation: during training it uses torch complex multiply (autograd-compatible); during inference it auto-dispatches to a fused CUDA kernel when available. The key property is that the dot product $q_i^T k_j$ depends only on the relative position $i - j$, not the absolute positions.
 
 **Critical for inference**: RoPE is applied **before** KV cache write, not after. If applied after caching, position encoding drift occurs because cached K/V would have stale rotation factors.
 
@@ -151,7 +151,7 @@ Three-layer separation (SGLang-inspired):
 - **ReqToTokenPool**: Index table `[req_idx, pos] → physical token slot`, shared across all layers.
 - **Allocator + PrefixCache**: Paged-mode slot allocation with ref-counting, LRU eviction, and hash-based prefix sharing.
 
-`PagePool` orchestrates all three. In contiguous mode (default), `req_to_token` is a trivial linear mapping. In paged mode, slots are allocated on demand with prefix caching support. Attention layers access buffers directly via `KVCache` dataclass — no methods, no abstraction.
+`PagePool` orchestrates all three. In contiguous mode (default), `req_to_token` is a trivial linear mapping. In paged mode, slots are allocated on demand with prefix caching support. `bind_tasks()` returns a `KVCache` dataclass with precomputed `page_table` and `decode_mask` fields (computed once per decode step, shared across all layers). Attention layers access buffers directly — no methods, no abstraction.
 
 ### Attention Backend
 
@@ -160,6 +160,8 @@ Attention computation is decoupled from the model via `AttentionBackend` ABC (`a
 - **`TorchNativeBackend`** (default): writes K/V to cache, gathers via `req_to_token` indirect indexing, calls `F.scaled_dot_product_attention`.
 - **`CudaBackend`**: decode path uses `attn_paged_decode` with `page_size=1` (the `req_to_token` table serves as the page table, each token slot is a single-token "page"); prefill path gathers K/V then calls `attn_prefill`. Falls back to `TorchNativeBackend` when kernel unavailable.
 
+Rotary embedding is applied via `apply_rotary_emb` in `astrai/extension/rotary_backend.py`, which auto-dispatches to the fused CUDA kernel (`rotary_emb.cu`) during inference or torch complex multiply during training (for autograd compatibility). Both attention backends share the same rotary dispatch.
+
 Backend selection is thread-safe via `contextvars`, mirroring `torch.nn.attention.sdpa_kernel`:
 
 ```python
@@ -228,4 +230,4 @@ total_steps          = (batches_per_replica // grad_accum_steps) * n_epoch
 
 This accounts for data-parallel sharding — each rank processes `1/nprocs` of the dataset.
 
-> Document Update Time: 2026-07-30
+> Document Update Time: 2026-07-31

docs/get-started.md

@@ -17,14 +17,14 @@ cd AstrAI
 # Basic install (pure PyTorch, no custom CUDA kernels)
 pip install -e .
 
-# With CUDA kernels (optional, for fused attention)
+# With CUDA kernels (optional, for fused attention and rotary embedding)
 # CSRC_KERNELS=true pip install -e . --no-build-isolation
 
 # With dev dependencies (pytest, ruff)
 # pip install -e ".[dev]"
 ```
 
-> **CUDA kernels** are opt-in. They are not built by default. When built, they can be activated via `with attn_backend(ATTN_BACKEND.CUDA):` for accelerated decode/prefill. You can skip them for normal usage.
+> **CUDA kernels** are opt-in. They are not built by default. When built, they can be activated via `with attn_backend(ATTN_BACKEND.CUDA):` for accelerated decode/prefill, and the fused rotary embedding kernel is auto-dispatched when available. You can skip them for normal usage.
 
 ## 2. Download Model Weights
 
@@ -232,4 +232,4 @@ docker compose up -d
 | System architecture | [Architecture](developer/architecture.md) |
 | Data pipeline internals | [Data Flow](developer/dataflow.md) |
 
-> Document Update Time: 2026-07-30
+> Document Update Time: 2026-07-31

docs/guides/inference.md

@@ -47,7 +47,10 @@ KVCache
   ├── req_to_token           [num_reqs, max_ctx_len]
   ├── req_pool_indices       [batch_size]
   ├── seq_lens               [batch_size]
-  └── out_cache_loc          [batch, seq_len] — write indices for this forward
+  ├── out_cache_loc          [batch, seq_len] — write indices for this forward
+  ├── max_len                int — max(seq_lens), avoids GPU sync in decode
+  ├── page_table             [batch, max_len] — precomputed gather indices for decode (None for prefill)
+  └── decode_mask            [batch, max_len] bool — precomputed position validity mask (None for single-batch decode)
 ```
 
 Attention layers do raw buffer indexing: `k_buffer[layer_id, out_cache_loc] = k` to write, `k_buffer[layer_id, indices]` to gather.
@@ -77,6 +80,15 @@ with attn_backend(ATTN_BACKEND.CUDA):
 
 Fallback: `CudaBackend` delegates to `TorchNativeBackend` when a CUDA kernel is not available.
 
+### Rotary Embedding Backend
+
+Rotary embedding is applied via `apply_rotary_emb` in `astrai/extension/rotary_backend.py`, which auto-dispatches:
+
+- **CUDA kernel** (`rotary_emb.cu`): fused cos/sin lookup + rotation in a single kernel, used when the kernel is available, input is on CUDA, and `torch.is_grad_enabled()` is `False` (inference mode)
+- **Torch fallback**: complex multiply path (`torch.view_as_complex` → `torch.complex` multiply → `torch.view_as_real`), used during training (supports autograd backward) or when the CUDA kernel is not available
+
+`RotaryEmbedding` stores `cos_table`/`sin_table` as f32 buffers and returns a `(cos, sin)` tuple from `forward()`. Both attention backends share the same rotary dispatch — it is backend-agnostic.
+
 ## Continuous Batching
 
 `InferenceScheduler` runs a daemon thread with a 4-phase loop:
@@ -183,6 +195,10 @@ Supports `stop_sequences` and streaming via `event: content_block_delta`.
 | `temperature` | float | 1.0 | Sampling temperature (> 0.0) |
 | `max_tokens` | Optional[int] | None | Max generation length |
 | `stream` | bool | False | Stream output |
+| `stop` | Optional[Union[str, List[str]]] | None | Stop sequences |
+| `frequency_penalty` | float | 0.0 | Frequency penalty |
+| `tools` | Optional[List[dict]] | None | Tool definitions for function calling |
+| `tool_choice` | Optional[str] | None | Tool selection mode |
 
 ### SSE Streaming Format
 
@@ -278,4 +294,4 @@ async for token in engine.generate_async("Hello", ...):    # -> AsyncGenerator[s
     print(token)
 ```
 
-> Document Update Time: 2026-07-30
+> Document Update Time: 2026-07-31

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 |

docs/guides/training.md

@@ -41,7 +41,7 @@ RoPE embeds position into Q/K vectors via complex rotation:
 
 $$ q_i = R_i W_q x_i, \quad k_j = R_j W_k x_j, \quad q_i^T k_j = x_i^T W_q^T R_{i-j} W_k x_j $$
 
-The complex rotation `freqs_cis` is pre-computed once (`cos, sin` pairs per position). `apply_rotary_emb` multiplies Q/K as complex numbers.
+`RotaryEmbedding` pre-computes `cos_table` and `sin_table` (f32, `[max_len, dim/2]`). `forward()` returns a `(cos, sin)` tuple indexed by `position_ids`. `apply_rotary_emb` applies the rotation: during training it uses torch complex multiply (autograd-compatible); during inference it auto-dispatches to a fused CUDA kernel when available.
 
 ## Training Loop
 
@@ -232,4 +232,4 @@ nohup python scripts/tools/train.py \
 
 Full parameter reference at [params.md](params.md).
 
-> Document Update Time: 2026-07-20
+> Document Update Time: 2026-07-31

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)