fix: resolve audited training, import, and serving bugs

- shard the Muon Newton-Schulz orthogonalization over the FSDP mesh instead of partial local slices
- import HF checkpoints faithfully: per-head RoPE permutation for q/k projections and qk-norm, qwen3, shared experts, and qk-norm before RoPE (changes numerics for existing use_qk_norm checkpoints)
- make preprocessing and resume self-contained: backfill realigned bucket keys by semantics (masks ones, rest zeros) and snapshot tokenizer files into every checkpoint
- keep RL consistent: sync the offline GRPO old_model each optimizer step and validate online strategies through a public one-off-rollout hook that leaves the replay cache untouched
- fix streaming serving: withhold partial tool-call prefixes with a stream-end flush, stream tool-call arguments from the raw source span, and terminate SSE frames with a blank line
- fix sampling semantics: capture logprobs before top-k/top-p mutate logits in place and detect greedy pipelines polymorphically instead of isinstance bookkeeping

astrai/optim/muon_adamw.py

@@ -1,9 +1,16 @@
 """Legacy Muon + AdamW combined optimizer."""
 
+from collections.abc import Mapping
 from typing import Any
 
 import torch
 from torch import Tensor, nn, optim
+from torch.distributed.tensor import DTensor, distribute_tensor
+from torch.optim._muon import (
+    _adjust_lr,
+    _single_tensor_muon,
+    _zeropower_via_newtonschulz,
+)
 
 from astrai.optim.composite import (
     OptimizerFactory,
@@ -14,6 +21,91 @@ from astrai.optim.composite import (
 )
 
 
+def _scalar_lr(lr: Any) -> float:
+    return lr.item() if isinstance(lr, Tensor) else lr
+
+
+def _sharded_orthogonalize(update: Tensor, group: Mapping) -> Tensor:
+    """Newton-Schulz for a sharded DTensor momentum update.
+
+    NS needs global matmuls, so gather the update to the full matrix,
+    orthogonalize it, and scatter the result back onto the update's
+    shard layout. ``full_tensor()`` returns the same gathered matrix on
+    every rank, so the scatter is a uniform collective.
+    """
+    full = update.full_tensor()
+    ortho = _zeropower_via_newtonschulz(
+        full, group["ns_coefficients"], group["ns_steps"], group["eps"]
+    )
+    return distribute_tensor(ortho, update.device_mesh, update.placements)
+
+
+class _ShardedMuon(optim.Muon):
+    """Muon that materializes sharded DTensor params around Newton-Schulz.
+
+    FSDP2 hands this optimizer dim-0 sharded DTensor parameters. The NS
+    iteration needs global matmuls: run it on the gathered full matrix,
+    then scatter the orthogonalized update back onto the parameter's
+    sharded layout so momentum buffers and weight decay stay sharded.
+    Without this, ``og @ og.T`` produces ``Partial(sum)`` DTensors that
+    downstream ``addmm`` calls consume without completing the reduction,
+    silently corrupting every update (measured 2e-4-9e-4 relative error
+    per step at world_size=2).
+
+    Plain (non-DTensor) params are routed through torch's own
+    ``_single_tensor_muon`` so unsharded runs stay bit-for-bit identical
+    to ``optim.Muon`` and this class carries only the DTensor delta.
+    Element-wise ops (momentum lerp, weight decay, the final ``add_``)
+    are DTensor-safe and run directly on the shards.
+    """
+
+    @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:
+            params: list[Tensor] = []
+            grads: list[Tensor] = []
+            bufs: list[Tensor] = []
+            self._init_group(group, params, grads, bufs)
+
+            plain, sharded = [], []
+            for param, grad, buf in zip(params, grads, bufs):
+                (sharded if isinstance(param, DTensor) else plain).append(
+                    (param, grad, buf)
+                )
+
+            if plain:
+                pp, gg, bb = (list(t) for t in zip(*plain))
+                _single_tensor_muon(
+                    pp,
+                    gg,
+                    bb,
+                    lr=group["lr"],
+                    weight_decay=group["weight_decay"],
+                    momentum=group["momentum"],
+                    nesterov=group["nesterov"],
+                    ns_coefficients=group["ns_coefficients"],
+                    ns_steps=group["ns_steps"],
+                    eps=group["eps"],
+                    adjust_lr_fn=group["adjust_lr_fn"],
+                    has_complex=False,
+                )
+
+            lr = _scalar_lr(group["lr"])
+            for param, grad, buf in sharded:
+                buf.lerp_(grad, 1 - group["momentum"])
+                update = grad.lerp(buf, group["momentum"]) if group["nesterov"] else buf
+
+                adjusted_lr = _adjust_lr(lr, group["adjust_lr_fn"], param.shape)
+                param.mul_(1 - lr * group["weight_decay"])
+                param.add_(_sharded_orthogonalize(update, group), alpha=-adjusted_lr)
+        return loss
+
+
 @OptimizerFactory.register("muon_adamw")
 class MuonAdamW(optim.Optimizer):
     """Combined Muon (matrix) + AdamW (non-matrix) optimizer."""
@@ -57,7 +149,7 @@ class MuonAdamW(optim.Optimizer):
             else:
                 other_params.append(param)
 
-        self.muon = optim.Muon(
+        self.muon = _ShardedMuon(
             matrix_params,
             lr=lr,
             weight_decay=weight_decay,