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/inference/network/openai.py

@@ -204,10 +204,31 @@ class OpenAIResponseBuilder(ResponseBuilder):
         )
 
     def format_stream_end(self, ctx: GenContext, stop: StopInfo) -> List[str]:
+        events: List[str] = []
+        if self._parser is not None:
+            for d in self._parser.finalize(stop.body):
+                if "content" in d:
+                    events.append(
+                        sse_event(
+                            {
+                                "id": self._resp_id,
+                                "object": "chat.completion.chunk",
+                                "created": ctx.created,
+                                "model": self._model,
+                                "choices": [
+                                    {
+                                        "index": 0,
+                                        "delta": {"content": d["content"]},
+                                        "finish_reason": None,
+                                    }
+                                ],
+                            }
+                        )
+                    )
         finish_reason = "stop"
         if self._parser is not None and self._parser.has_tool_calls:
             finish_reason = "tool_calls"
-        return [
+        return events + [
             sse_event(
                 {
                     "id": self._resp_id,

astrai/inference/network/protocol.py

@@ -20,8 +20,10 @@ def sse_event(data: Dict[str, Any], event: Optional[str] = None) -> str:
     if event:
         lines.append(f"event: {event}")
     lines.append(f"data: {json.dumps(data, ensure_ascii=False)}")
-    lines.append("")
-    return "\n".join(lines)
+    # The SSE spec dispatches an event only at a blank line, so the frame
+    # must end with "\n\n" (a single trailing newline keeps clients waiting
+    # and concatenates consecutive events into one corrupt payload).
+    return "\n".join(lines) + "\n\n"
 
 
 def sse_done() -> str:

astrai/inference/network/tool_parser.py

@@ -69,6 +69,10 @@ class BaseToolParser(ABC):
     def has_tool_calls(self) -> bool:
         """True if the parser detected at least one tool call in the stream."""
 
+    def finalize(self, body: str) -> List[Dict]:
+        """Flush parser state once generation ended. Default: nothing."""
+        return []
+
 
 class ToolParserFactory(BaseFactory["BaseToolParser"]):
     pass
@@ -108,6 +112,30 @@ def _scan_json(text: str, start: int = 0):
     return len(text), False
 
 
+def _raw_arguments_span(json_str: str) -> Optional[str]:
+    """Extract the raw text span of the ``arguments`` value, if possible.
+
+    Streaming diffs emit the *source* text of ``arguments``; using the
+    same span here keeps the completed arguments a prefix-extension of
+    what was already streamed (``json.dumps`` would re-quote and
+    re-space the value and corrupt the concatenated result).
+    """
+    m = re.search(r'"arguments"\s*:\s*', json_str)
+    if not m:
+        return None
+    rest = json_str[m.end() :]
+    if rest[:1] in ("{", "["):
+        end, ok = _scan_json(rest, 0)
+        if ok:
+            return rest[1 : end - 1]
+        return None
+    str_match = re.match(r'"(?:[^"\\]|\\.)*"', rest)
+    if str_match:
+        return str_match.group(0)
+    bare = re.match(r"[^,}\s][^,}]*", rest)
+    return bare.group(0).rstrip() if bare else None
+
+
 def _parse_tool_call_json(json_str: str, complete: bool):
     """Extract *name* and *arguments* from a tool-call JSON string.
 
@@ -122,14 +150,23 @@ def _parse_tool_call_json(json_str: str, complete: bool):
         if not isinstance(name, str) or not name:
             return None, "", False
         args = obj.get("arguments")
+        raw = _raw_arguments_span(json_str)
         if isinstance(args, dict):
             if not args:
                 args = ""
             else:
-                args = json.dumps(args, ensure_ascii=False)
-                args = args[1:-1].rstrip()
+                # Prefer the source span: it matches what streaming
+                # already emitted (prefix-consistent completion).
+                args = (
+                    raw
+                    if raw is not None
+                    else json.dumps(args, ensure_ascii=False)[1:-1].rstrip()
+                )
         elif isinstance(args, list):
-            args = json.dumps(args, ensure_ascii=False) if args else ""
+            if not args:
+                args = ""
+            else:
+                args = raw if raw is not None else json.dumps(args, ensure_ascii=False)
         elif isinstance(args, str):
             pass
         else:
@@ -306,10 +343,52 @@ class SimpleJsonToolParser(BaseToolParser):
         return deltas
 
     def _emit_plain_content(self, body: str, deltas: List[Dict]) -> List[Dict]:
-        new_content = body[self._emitted_content_len :]
-        if new_content:
+        safe_end = self._safe_content_end(body)
+        if safe_end > self._emitted_content_len:
+            deltas.append({"content": body[self._emitted_content_len : safe_end]})
+            self._emitted_content_len = safe_end
+        return deltas
+
+    _POSSIBLE_NAME_PREFIX_RE = re.compile(r'^\s*"?(?:n|na|nam|name)?"?\s*:?\s*$')
+
+    @classmethod
+    def _safe_content_end(cls, body: str) -> int:
+        """End of *body* that is safe to emit as plain content.
+
+        A trailing unclosed ``{`` whose remainder could still grow into
+        ``{"name": ...`` (the name-prefix itself, or any further ``{``
+        opened after it) is withheld, otherwise a partial ``{"na`` prefix
+        would leak into user-visible content and never be retracted.
+        Anything withheld is flushed by :meth:`finalize` when generation
+        ends without a tool call, so plain text is never lost.
+        """
+        pos = 0
+        while True:
+            brace = body.find("{", pos)
+            if brace == -1:
+                return len(body)
+            end, complete = _scan_json(body, brace)
+            if complete:
+                pos = end
+                continue
+            tail = body[brace + 1 :]
+            if cls._POSSIBLE_NAME_PREFIX_RE.match(tail) or "{" in tail:
+                return brace
+            return len(body)
+
+    def finalize(self, body: str) -> List[Dict]:
+        """Flush content withheld as a possible tool-call prefix.
+
+        Called once generation completed: if no tool call materialised,
+        emit the withheld remainder so streamed content matches the
+        non-streaming response.
+        """
+        if self._has_tool_calls:
+            return []
+        deltas: List[Dict] = []
+        if len(body) > self._emitted_content_len:
+            deltas.append({"content": body[self._emitted_content_len :]})
             self._emitted_content_len = len(body)
-            deltas.append({"content": new_content})
         return deltas
 
     # -------------------------------------------------------- complete

astrai/inference/runtime/sample.py

@@ -42,6 +42,22 @@ class BaseSamplingStrategy(ABC):
         """
         raise NotImplementedError
 
+    @property
+    def preserves_argmax(self) -> bool:
+        """Whether ``apply`` never moves the argmax token.
+
+        Conservative default: strategies must opt in. The greedy
+        short-circuit in :class:`SamplingPipeline` asks this
+        polymorphically, so a new strategy that can move the argmax
+        automatically disables it — no isinstance bookkeeping.
+        """
+        return False
+
+    @property
+    def is_greedy(self) -> bool:
+        """Whether this strategy collapses sampling onto the argmax token."""
+        return False
+
 
 class TemperatureStrategy(BaseSamplingStrategy):
     """Divides logits by temperature to control randomness.
@@ -53,6 +69,22 @@ class TemperatureStrategy(BaseSamplingStrategy):
     def __init__(self, temperature: Union[float, Tensor] = 1.0):
         self.temperature = temperature
 
+    @staticmethod
+    def is_greedy_temperature(temperature: Union[float, Tensor]) -> bool:
+        if isinstance(temperature, Tensor):
+            return bool((temperature == 0).all())
+        return temperature == 0
+
+    @property
+    def is_greedy(self) -> bool:
+        return self.is_greedy_temperature(self.temperature)
+
+    @property
+    def preserves_argmax(self) -> bool:
+        # Scaling by a positive constant (1/t, clamped away from zero)
+        # preserves logit order; t=0 degenerates onto the argmax itself.
+        return True
+
     def apply(
         self,
         logits: Tensor,
@@ -78,6 +110,11 @@ class TopKStrategy(BaseSamplingStrategy):
         top_k: Scalar or ``[batch]`` tensor (0 disables).
     """
 
+    @property
+    def preserves_argmax(self) -> bool:
+        # The argmax token always ranks first, so any k >= 1 keeps it.
+        return True
+
     def __init__(self, top_k: Union[int, Tensor] = 0):
         self.top_k = top_k
 
@@ -121,6 +158,11 @@ class TopPStrategy(BaseSamplingStrategy):
         top_p: Scalar or ``[batch]`` tensor (1.0 disables).
     """
 
+    @property
+    def preserves_argmax(self) -> bool:
+        # Nucleus filtering always keeps the highest-probability token.
+        return True
+
     def __init__(self, top_p: Union[float, Tensor] = 1.0):
         self.top_p = top_p
 
@@ -250,6 +292,22 @@ class SamplingPipeline(BaseSamplingStrategy):
     def __init__(self, strategies: List[BaseSamplingStrategy]):
         self.strategies = strategies
 
+    @property
+    def preserves_argmax(self) -> bool:
+        # A composite preserves the argmax iff every stage does.
+        return all(s.preserves_argmax for s in self.strategies)
+
+    @property
+    def is_greedy(self) -> bool:
+        """Whether sampling always yields the argmax of the raw logits.
+
+        True iff some stage forces greedy and no stage can move the
+        argmax before or after it. Both facts are declared
+        polymorphically by each strategy, so composing in a new strategy
+        type (or a nested pipeline) updates this automatically.
+        """
+        return any(s.is_greedy for s in self.strategies) and self.preserves_argmax
+
     def apply(
         self,
         logits: Tensor,
@@ -261,12 +319,6 @@ class SamplingPipeline(BaseSamplingStrategy):
             logits = strategy.apply(logits, filter_value, input_ids, input_mask)
         return logits
 
-    @staticmethod
-    def _is_greedy(temperature: Union[float, Tensor]) -> bool:
-        if isinstance(temperature, Tensor):
-            return bool((temperature == 0).all())
-        return temperature == 0
-
     @torch.inference_mode()
     def sample(
         self,
@@ -294,7 +346,7 @@ class SamplingPipeline(BaseSamplingStrategy):
             Sampled token IDs ``[batch]``, or — when ``return_logprobs``
             is ``True`` — a ``(token_ids, chosen_logprobs)`` tuple.
         """
-        if self._is_greedy_pipeline():
+        if self.is_greedy:
             tokens = logits.argmax(dim=-1)
             if not return_logprobs:
                 return tokens
@@ -302,6 +354,13 @@ class SamplingPipeline(BaseSamplingStrategy):
             chosen = torch.gather(log_probs, -1, tokens.unsqueeze(-1)).squeeze(-1)
             return tokens, chosen
 
+        # Capture the raw distribution before the strategy pipeline runs:
+        # top-k/top-p mutate the logits tensor in place, so computing this
+        # after ``apply`` would read the filtered distribution instead of
+        # the raw model distribution the caller documented.
+        if return_logprobs:
+            raw_log_probs = torch.log_softmax(logits.float(), dim=-1)
+
         transformed = self.apply(logits, filter_value, input_ids, input_mask)
         tokens = torch.multinomial(
             torch.softmax(transformed, dim=-1), num_samples=1
@@ -313,28 +372,9 @@ class SamplingPipeline(BaseSamplingStrategy):
         # logprobs recorded for online RL must live in the same
         # distribution the trainer differentiates, not the
         # temperature/top-p filtered one tokens were drawn from.
-        log_probs = torch.log_softmax(logits.float(), dim=-1)
-        chosen = torch.gather(log_probs, -1, tokens.unsqueeze(-1)).squeeze(-1)
+        chosen = torch.gather(raw_log_probs, -1, tokens.unsqueeze(-1)).squeeze(-1)
         return tokens, chosen
 
-    def _is_greedy_pipeline(self) -> bool:
-        """True if sampling reduces to argmax over the raw logits.
-
-        A greedy temperature with only top-k/top-p strategies does: the
-        filters always keep the argmax token.  A frequency penalty can
-        change the argmax, so those pipelines must run the full
-        transformation even at ``temperature=0``.
-        """
-        if not self.strategies:
-            return False
-        first = self.strategies[0]
-        if not (
-            isinstance(first, TemperatureStrategy)
-            and self._is_greedy(first.temperature)
-        ):
-            return False
-        return not any(isinstance(s, FrequencyPenaltyStrategy) for s in self.strategies)
-
 
 @torch.inference_mode()
 def sample(

astrai/model/components/attention.py

@@ -69,11 +69,15 @@ class GQA(nn.Module):
         q = self._split_heads(self.q_proj(x), self.n_heads)
         k = self._split_heads(self.k_proj(x), self.n_kv_heads)
         v = self._split_heads(self.v_proj(x), self.n_kv_heads)
-        q, k = apply_rotary_emb(q, rotary_emb), apply_rotary_emb(k, rotary_emb)
 
+        # Match the HuggingFace convention (Qwen2/Gemma): normalize Q/K
+        # before RoPE. RMSNorm's per-channel gain does not commute with
+        # the rotation, so the order changes numerics.
         if self.use_qk_norm:
             q, k = self.q_norm(q), self.k_norm(k)
 
+        q, k = apply_rotary_emb(q, rotary_emb), apply_rotary_emb(k, rotary_emb)
+
         sdqa_out = attention(
             q, k, v, kv_cache, self.layer_id, attn_mask, is_causal, fwd
         ).reshape(*x.shape[:-1], self.dim)

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,

astrai/parallel/executor.py

@@ -61,12 +61,13 @@ def broadcast_state_dict(
 
     rank = dist.get_rank()
 
-    # Broadcast metadata (keys, shapes, dtypes, device) so non-src ranks
-    # can allocate matching empty tensors on the correct device.
+    # Broadcast metadata (keys, shapes, dtypes, device kind) so non-src
+    # ranks can allocate matching empty tensors on their own device.
     if rank == src:
-        device = next(iter(state_dict.values())).device
+        first = next(iter(state_dict.values()), None)
+        on_cuda = first is not None and first.is_cuda
         metadata = [
-            (k, tuple(v.shape), v.dtype, str(device)) for k, v in state_dict.items()
+            (k, tuple(v.shape), v.dtype, on_cuda) for k, v in state_dict.items()
         ]
     else:
         metadata = None
@@ -75,10 +76,18 @@ def broadcast_state_dict(
     metadata = metadata_list[0]
 
     # Non-src ranks allocate empty tensors with the broadcasted metadata.
+    # The allocation must sit on *this* rank's local device: process groups
+    # are pinned per-rank via ``device_id=``, so NCCL rejects a tensor
+    # allocated on the src rank's device (and the cross-GPU allocation
+    # would be wrong here anyway).
     if rank != src:
+        device = (
+            torch.device("cuda", torch.cuda.current_device())
+            if any(m[3] for m in metadata)
+            else torch.device("cpu")
+        )
         state_dict = {
-            k: torch.empty(s, dtype=d, device=torch.device(dev))
-            for k, s, d, dev in metadata
+            k: torch.empty(s, dtype=d, device=device) for k, s, d, _ in metadata
         }
 
     # Broadcast each tensor in-place.
@@ -400,17 +409,16 @@ class FSDPExecutor(BaseExecutor):
         if not self.use_distributed:
             return super().clip_grad_norm(model, max_norm)
 
-        # FSDP params are DTensors (sharded across ranks).
-        # torch.nn.utils.clip_grad_norm_ computes LOCAL norm per rank,
-        # so we must all-reduce to get the global norm before clipping.
-        local_norm = torch.nn.utils.get_total_norm(
+        # FSDP params are DTensors (sharded across ranks), and
+        # get_total_norm reduces them to a globally-reduced replicated
+        # DTensor — no extra all-reduce is needed. Manually summing the
+        # local slice again would inflate the norm by sqrt(world_size)
+        # and over-clip every step.
+        total_norm = torch.nn.utils.get_total_norm(
             [p.grad for p in model.parameters() if p.grad is not None],
         )
-        if isinstance(local_norm, DTensor):
-            local_norm = local_norm.to_local()
-        total_norm_sq = local_norm**2
-        dist.all_reduce(total_norm_sq, op=dist.ReduceOp.SUM)
-        total_norm = total_norm_sq.sqrt()
+        if isinstance(total_norm, DTensor):
+            total_norm = total_norm.full_tensor()
 
         clip_coef = max_norm / (total_norm + 1e-6)
         clip_coef_clamped = torch.clamp(clip_coef, max=1.0)

astrai/preprocessing/pipeline.py

@@ -146,11 +146,17 @@ class Pipeline:
 
     @staticmethod
     def _align_bucket(bucket: dict, result: dict, ids: list):
-        """Pad previously-accumulated keys that are missing from *result*."""
+        """Pad previously-accumulated keys that are missing from *result*.
+
+        Builders omit all-ones masks (``loss_mask`` / ``*_mask``) to save
+        space, so an omitted mask means "train on every token" and is
+        back-filled with ones; every other missing key pads with zeros.
+        """
         for key in list(bucket.keys()):
             if key in result:
                 continue
-            bucket[key].append([0] * len(ids))
+            fill = 1 if key == "loss_mask" or key.endswith("_mask") else 0
+            bucket[key].append([fill] * len(ids))
 
     def _iter_items(self):
         for path in self.paths:

astrai/serialization/hf_adapter.py

@@ -36,6 +36,7 @@ HF_MODEL_TYPES = frozenset(
         "mixtral",
         "qwen2",
         "qwen2_moe",
+        "qwen3",
         "gemma",
         "gemma2",
         "phi3",
@@ -58,13 +59,29 @@ _MOE_EXPERTS = re.compile(
     r"^model\.layers\.(\d+)\.mlp\.experts\.(\d+)\.(gate|up|down)_proj\.(weight|bias)$"
 )
 _MOE_SHARED = re.compile(
-    r"^model\.layers\.(\d+)\.mlp\.shared_expert(?:s)?\.(\d+)\."
+    r"^model\.layers\.(\d+)\.mlp\.shared_expert(?:s)?(?:\.(\d+))?\."
     r"(gate|up|down)_proj\.(weight|bias)$"
 )
 
 _ASTR_PREFIXES = ("embed_tokens.", "layers.", "norm.", "lm_head.")
 
 
+def _half_to_interleaved(head_dim: int) -> torch.Tensor:
+    """Row permutation converting HF half-split RoPE coordinates to
+    AstrAI interleaved coordinates.
+
+    HF rotate_half pairs channels ``(i, i + head_dim/2)``; AstrAI pairs
+    ``(2i, 2i + 1)``. Both use frequency ``i`` for the pair, so AstrAI
+    channel ``2i`` takes the HF value at channel ``i`` and AstrAI
+    ``2i + 1`` takes HF ``i + head_dim/2``.
+    """
+    half = head_dim // 2
+    perm = torch.empty(head_dim, dtype=torch.long)
+    perm[0::2] = torch.arange(half)
+    perm[1::2] = torch.arange(half, head_dim)
+    return perm
+
+
 def looks_like_hf_state_dict(state_dict: Mapping[str, Any]) -> bool:
     """Return True if *state_dict* uses HuggingFace key names."""
     return any(
@@ -168,8 +185,11 @@ def convert_hf_config(raw: Dict[str, Any]) -> Dict[str, Any]:
             cfg["n_activated_experts"] = raw["n_activated_experts"]
         if "n_shared_experts" in raw:
             cfg["n_shared_experts"] = raw["n_shared_experts"]
+        elif raw.get("shared_expert_intermediate_size"):
+            # Qwen2-MoE exposes a single un-indexed shared expert.
+            cfg["n_shared_experts"] = 1
         else:
-            # Mixtral has no shared experts; AstrAI defaults to one.
+            # Mixtral has no shared experts.
             cfg["n_shared_experts"] = 0
         if cfg.get("moe_intermediate_size") is None and "intermediate_size" in raw:
             # MoE configs store the per-expert FFN size in intermediate_size.
@@ -201,6 +221,14 @@ def convert_hf_weights(
             )
 
     ffn_type = getattr(config, "ffn_type", "mlp")
+    permute_rope = getattr(config, "attn_type", "gqa") != "mla"
+    head_dim = None
+    if permute_rope:
+        head_dim = config.hidden_size // config.num_attention_heads
+        if head_dim % 2 != 0:
+            raise ValueError(
+                f"head_dim={head_dim} is odd; rotary permutation requires even"
+            )
     converted: Dict[str, torch.Tensor] = {}
     skipped: list[str] = []
     for key, tensor in state_dict.items():
@@ -223,8 +251,9 @@ def convert_hf_weights(
                 else:
                     m = _MOE_SHARED.match(key)
                     if m:
+                        shared_idx = m.group(2) if m.group(2) is not None else "0"
                         new_key = (
-                            f"layers.{m.group(1)}.mlp.shared_experts.{m.group(2)}."
+                            f"layers.{m.group(1)}.mlp.shared_experts.{shared_idx}."
                             f"{m.group(3)}.{m.group(4)}"
                         )
             if new_key is None:
@@ -242,10 +271,31 @@ def convert_hf_weights(
                 new_key = (
                     f"layers.{m.group(1)}.attention.{m.group(2)}_proj.{m.group(3)}"
                 )
+                if permute_rope and m.group(2) in ("q", "k"):
+                    rows = tensor.shape[0]
+                    if rows % head_dim != 0:
+                        raise ValueError(
+                            f"{key}: {rows} output rows not divisible by "
+                            f"head_dim={head_dim}"
+                        )
+                    base = _half_to_interleaved(head_dim).to(tensor.device)
+                    blocks = (
+                        torch.arange(rows // head_dim, device=tensor.device) * head_dim
+                    )
+                    perm = (blocks[:, None] + base[None, :]).flatten()
+                    tensor = tensor.index_select(0, perm)
             elif (m := _Q_NORM.match(key)) is not None:
                 new_key = f"layers.{m.group(1)}.attention.q_norm.weight"
+                if permute_rope and tensor.shape[0] == head_dim:
+                    tensor = tensor.index_select(
+                        0, _half_to_interleaved(head_dim).to(tensor.device)
+                    )
             elif (m := _K_NORM.match(key)) is not None:
                 new_key = f"layers.{m.group(1)}.attention.k_norm.weight"
+                if permute_rope and tensor.shape[0] == head_dim:
+                    tensor = tensor.index_select(
+                        0, _half_to_interleaved(head_dim).to(tensor.device)
+                    )
             elif (m := _INPUT_NORM.match(key)) is not None:
                 new_key = f"layers.{m.group(1)}.input_norm.weight"
             elif (m := _POST_NORM.match(key)) is not None:

astrai/trainer/rollout.py

@@ -551,3 +551,16 @@ class RolloutRunner:
         # cache publication. Reward scoring itself intentionally remains
         # outside the policy lock because it may call an external service.
         return self.generator.with_policy_snapshot(commit)
+
+    def evaluate(self, batch: Dict) -> RolloutResult:
+        """One-off rollout + scoring that leaves the replay cache untouched.
+
+        Used by validation on online strategies: the training cache, its
+        cadence counter, and the cache key stay intact, so evaluation
+        prompts never disturb the rollout replay schedule.
+        """
+        raw = self.generator.generate(batch)
+        self._validate_policy_version(raw)
+        scored = self._score(raw)
+        self._validate_policy_version(scored)
+        return scored

astrai/trainer/strategy.py

@@ -1,7 +1,7 @@
 """Training strategy implementations with factory pattern."""
 
 from abc import ABC
-from typing import Callable, Dict, List, Optional, TypedDict, Union
+from typing import Any, Callable, Dict, List, Optional, TypedDict, Union
 
 import torch
 import torch.nn as nn
@@ -202,6 +202,21 @@ class BaseStrategy(ABC):
     def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput:
         return self._normalize_output(self.compute_loss(batch))
 
+    def validate_online(self, batch: Dict[str, Any]) -> Optional[LossOutput]:
+        """Validate one batch through a one-off rollout.
+
+        Online strategies with an injected rollout runner evaluate a
+        fresh, throw-away rollout so the training replay cache and its
+        cadence stay untouched. Returns ``None`` when no runner is
+        configured (offline mode); callers then fall back to
+        ``strategy(batch)``.
+        """
+        if self._rollout_runner is None:
+            return None
+        result = self._rollout_runner.evaluate(batch)
+        prepared = self.prepare_from_rollout(result)
+        return self.compute_loss_output(prepared)
+
     def _loss_output(
         self,
         task_loss: Tensor,
@@ -595,6 +610,19 @@ class GRPOStrategy(BaseStrategy):
         if state_dict is not None:
             self.old_model.load_state_dict(state_dict)
 
+    def optimizer_step(self, optimizer: Optimizer):
+        """Step the optimizer, then refresh the offline behaviour policy.
+
+        Without this sync the frozen ``old_model`` drifts away from the
+        training policy, so the PPO ratio degenerates and clipping shuts
+        learning down. Online GRPO passes ``logprobs_old`` instead and
+        runs with ``old_model=None``, skipping the sync.
+        """
+        result = super().optimizer_step(optimizer)
+        if self.old_model is not None:
+            self.sync_old_model()
+        return result
+
     def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput:
         batch = move_to_device(batch, self.device)
         prompts = batch["prompts"]

astrai/trainer/train_callback.py

@@ -1,6 +1,7 @@
 import json
 import logging
 import os
+import shutil
 import sys
 import time
 from functools import partial
@@ -29,6 +30,32 @@ from astrai.trainer.train_context import TrainContext
 
 logger = logging.getLogger(__name__)
 
+_TOKENIZER_FILES = (
+    "tokenizer.json",
+    "tokenizer_config.json",
+    "special_tokens_map.json",
+)
+
+
+def _copy_tokenizer_files(param_path: Optional[str], save_path: str):
+    """Snapshot tokenizer files into the checkpoint directory.
+
+    ``param_path`` is the launch model directory (or, on resume, a
+    previous self-contained checkpoint), so the copy makes every
+    checkpoint independently resumable for online training, which
+    loads its tokenizer from ``param_path``.
+    """
+    if not param_path:
+        return
+    for name in _TOKENIZER_FILES:
+        src = os.path.join(param_path, name)
+        dst = os.path.join(save_path, name)
+        if not os.path.isfile(src) or (
+            os.path.isfile(dst) and os.path.samefile(src, dst)
+        ):
+            continue
+        shutil.copy2(src, dst)
+
 
 @runtime_checkable
 class TrainCallback(Protocol):
@@ -176,6 +203,7 @@ class CheckpointCallback(TrainCallback):
                     meta=meta,
                 )
                 context.checkpoint.save(save_path)
+                _copy_tokenizer_files(context.param_path, save_path)
         self.last_ckpt_step = context.optimizer_step
 
     def after_optimizer_step(self, context: TrainContext):
@@ -325,7 +353,12 @@ class MetricCallback(TrainCallback):
 
         with torch.no_grad():
             for batch in context.val_dataloader:
-                loss_output = context.strategy(batch)
+                # Online strategies evaluate a one-off rollout (leaving
+                # the replay cache untouched) via the public hook; None
+                # means offline — validate the batch directly.
+                loss_output = context.strategy.validate_online(batch)
+                if loss_output is None:
+                    loss_output = context.strategy(batch)
                 total_loss += loss_output["loss"].item()
                 num_batches += 1
 

astrai/trainer/train_context.py

@@ -59,6 +59,7 @@ class TrainContext:
     world_size: int = field(default=1)
     rank: int = field(default=0)
     kwargs: Dict[str, Any] = field(default_factory=dict)
+    param_path: Optional[str] = field(default=None)
 
     _stop_event: threading.Event = field(default_factory=threading.Event)
 
@@ -180,6 +181,7 @@ class TrainContextBuilder:
             epoch=state.epoch,
             consumed_samples=state.consumed_samples,
             checkpoint=state.checkpoint,
+            param_path=self._param_path,
         )
 
     def _prepare_model(