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("")
+    # The SSE spec dispatches an event only at a blank line, so the frame
-    return "\n".join(lines)
+    # 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)
+                # Prefer the source span: it matches what streaming
-                args = args[1:-1].rstrip()
+                # 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 :]
+        safe_end = self._safe_content_end(body)
-        if new_content:
+        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(raw_log_probs, -1, tokens.unsqueeze(-1)).squeeze(-1)
-        chosen = torch.gather(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
+    # Broadcast metadata (keys, shapes, dtypes, device kind) so non-src
-    # can allocate matching empty tensors on the correct device.
+    # 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))
+            k: torch.empty(s, dtype=d, device=device) for k, s, d, _ in metadata
-            for k, s, d, dev 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).
+        # FSDP params are DTensors (sharded across ranks), and
-        # torch.nn.utils.clip_grad_norm_ computes LOCAL norm per rank,
+        # get_total_norm reduces them to a globally-reduced replicated
-        # so we must all-reduce to get the global norm before clipping.
+        # DTensor — no extra all-reduce is needed. Manually summing the
-        local_norm = torch.nn.utils.get_total_norm(
+        # 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):
+        if isinstance(total_norm, DTensor):
-            local_norm = local_norm.to_local()
+            total_norm = total_norm.full_tensor()
-        total_norm_sq = local_norm**2
-        dist.all_reduce(total_norm_sq, op=dist.ReduceOp.SUM)
-        total_norm = total_norm_sq.sqrt()
 
         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(

tests/data/test_pipeline_align.py

@@ -0,0 +1,26 @@
+"""Tests for preprocessing pipeline bucket alignment."""
+
+from astrai.preprocessing.pipeline import Pipeline
+
+
+def test_align_bucket_backfills_missing_mask_with_ones():
+    bucket = {
+        "sequence": [[1, 2], [3, 4]],
+        "loss_mask": [[0, 1]],
+        "chosen_mask": [[1]],
+        "position_ids": [[0, 1]],
+    }
+    result = {"sequence": [5, 6, 7]}
+    Pipeline._align_bucket(bucket, result, [5, 6, 7])
+    assert bucket["loss_mask"][-1] == [1, 1, 1]
+    assert bucket["chosen_mask"][-1] == [1, 1, 1]
+    assert bucket["position_ids"][-1] == [0, 0, 0]
+    assert bucket["sequence"] == [[1, 2], [3, 4]]
+
+
+def test_align_bucket_keeps_present_keys():
+    bucket = {"sequence": [[1, 2]], "loss_mask": [[0, 1]]}
+    result = {"sequence": [9], "loss_mask": [1]}
+    Pipeline._align_bucket(bucket, result, [9])
+    assert bucket["loss_mask"] == [[0, 1]]
+    assert bucket["sequence"] == [[1, 2]]

tests/inference/test_sample.py

@@ -3,6 +3,7 @@
 import torch
 
 from astrai.inference.runtime.sample import (
+    BaseSamplingStrategy,
     FrequencyPenaltyStrategy,
     SamplingPipeline,
     TemperatureStrategy,
@@ -295,3 +296,45 @@ def test_greedy_respects_frequency_penalty():
     )
     # Token 0 saw four occurrences: 5 - 2*4 < 4, so the argmax flips.
     assert penalized.tolist() == [1]
+
+
+class _ArgmaxMovingStrategy(BaseSamplingStrategy):
+    """Custom strategy that can move the argmax — must disable greedy."""
+
+    def apply(
+        self, logits, filter_value=-float("inf"), input_ids=None, input_mask=None
+    ):
+        return torch.roll(logits, shifts=1, dims=-1)
+
+
+def test_greedy_detection_is_polymorphic():
+    """Greedy detection asks strategies polymorphically, no isinstance."""
+    base = [TemperatureStrategy(0.0), TopKStrategy(50), TopPStrategy(0.9)]
+    assert SamplingPipeline(list(base)).is_greedy is True
+    assert SamplingPipeline(base + [FrequencyPenaltyStrategy(0.5)]).is_greedy is False
+
+    # A custom argmax-moving strategy disables greedy even though the
+    # pipeline contains a greedy temperature — this is what isinstance
+    # bookkeeping in the old implementation could not see.
+    assert SamplingPipeline(base + [_ArgmaxMovingStrategy()]).is_greedy is False
+
+
+def test_greedy_detection_position_independent():
+    """Greedy temperature anywhere in the pipeline is detected."""
+    pipeline = SamplingPipeline([TopKStrategy(50), TemperatureStrategy(0.0)])
+    assert pipeline.is_greedy is True
+
+
+def test_greedy_detection_composes_across_nested_pipelines():
+    """A nested pipeline participates through the same interface."""
+    inner = SamplingPipeline([TemperatureStrategy(0.0), TopKStrategy(20)])
+    assert inner.is_greedy is True
+    assert SamplingPipeline([TopPStrategy(0.9), inner]).is_greedy is True
+    assert SamplingPipeline([inner, FrequencyPenaltyStrategy(0.5)]).is_greedy is False
+
+
+def test_nongreedy_temperature_is_not_greedy():
+    pipeline = SamplingPipeline(
+        [TemperatureStrategy(0.7), TopKStrategy(0), TopPStrategy(1.0)]
+    )
+    assert pipeline.is_greedy is False

tests/inference/test_tool_parser.py

@@ -565,3 +565,40 @@ def test_parser_uses_token_ids_for_detection():
     parser = TokenIdParser()
     parser.feed("hello", current_token_ids=[1, 999, 3])
     assert parser.has_tool_calls
+
+
+def test_streaming_partial_name_prefix_never_leaks_into_content():
+    parser = SimpleJsonToolParser()
+    parts = ["Hello ", '{"', '{"n', '{"na', '{"name"']
+    emitted = []
+    body = ""
+    for part in parts:
+        body += part
+        for d in parser.feed(body):
+            if "content" in d:
+                emitted.append(d["content"])
+    assert "".join(emitted) == "Hello "
+
+
+def test_finalize_flushes_withheld_plain_json_content():
+    parser = SimpleJsonToolParser()
+    text = 'Answer: {"price": 1}'
+    deltas = parser.feed(text)
+    streamed = "".join(d["content"] for d in deltas if "content" in d)
+    flushed = parser.finalize(text)
+    joined = streamed + "".join(d["content"] for d in flushed if "content" in d)
+    assert joined == text
+    assert not parser.has_tool_calls
+    assert parser.finalize(text) == []
+
+
+def test_streaming_args_concat_matches_parse_complete():
+    parser = SimpleJsonToolParser()
+    # Compact spacing: json.dumps would re-space this and desync the
+    # streamed arguments diff.
+    text = '{"name": "get_weather","arguments": {"city":"Beijing","unit":"c"}}'
+    _, args_chunks = _simulate_streaming(parser, text)
+    streamed = "".join(args_chunks)
+    completed = parser.parse_complete(text)["tool_calls"][0]["function"]["arguments"]
+    assert streamed == completed
+    assert streamed == '"city":"Beijing","unit":"c"'

tests/serialization/test_hf_adapter.py

@@ -15,6 +15,7 @@ from astrai.serialization import (
     looks_like_hf_state_dict,
     save_model,
 )
+from astrai.serialization.hf_adapter import _half_to_interleaved
 from tests.helpers import assert_state_dicts_equal, make_tiny_config
 
 LLAMA_RAW = {
@@ -47,10 +48,38 @@ MOE_RAW = {
 }
 
 
-def to_hf_keys(state_dict):
+def to_hf_keys(state_dict, head_dim=None):
-    """Rename AstrAI state dict keys to HuggingFace LLaMA-style names."""
+    """Rename AstrAI state dict keys to HuggingFace LLaMA-style names.
+
+    When *head_dim* is given, q/k projections and q/k norm weights are
+    also converted from AstrAI interleaved RoPE coordinates to the HF
+    half-split (rotate_half) convention, so the produced state dict is a
+    faithful HF-layout checkpoint.
+    """
     out = {}
     for key, tensor in state_dict.items():
+        if head_dim is not None:
+            name = key.split(".")
+            is_qk_proj = (
+                len(name) >= 4
+                and name[2] == "attention"
+                and name[3] in ("q_proj", "k_proj")
+            )
+            is_qk_norm = (
+                len(name) >= 4
+                and name[2] == "attention"
+                and name[3] in ("q_norm", "k_norm")
+                and name[4] == "weight"
+            )
+            if is_qk_proj or is_qk_norm:
+                inv = torch.argsort(_half_to_interleaved(head_dim))
+                rows = tensor.shape[0]
+                if rows > head_dim:
+                    blocks = torch.arange(rows // head_dim) * head_dim
+                    idx = (blocks[:, None] + inv[None, :]).flatten()
+                else:
+                    idx = inv
+                tensor = tensor.index_select(0, idx)
         if key == "embed_tokens.weight":
             out["model.embed_tokens.weight"] = tensor
         elif key == "norm.weight":
@@ -141,7 +170,9 @@ def test_adapt_config_passthrough():
 def test_convert_hf_weights_dense_roundtrip():
     cfg = make_tiny_config()
     model = AutoRegressiveLM(cfg)
-    converted = convert_hf_weights(to_hf_keys(model.state_dict()), cfg)
+    converted = convert_hf_weights(
+        to_hf_keys(model.state_dict(), cfg.hidden_size // cfg.num_attention_heads), cfg
+    )
     assert_state_dicts_equal(converted, model.state_dict())
 
 
@@ -157,7 +188,10 @@ def test_convert_hf_weights_moe_roundtrip():
     model = AutoRegressiveLM(cfg)
     hf_raw = convert_hf_config(MOE_RAW)
     hf_cfg = ConfigFactory.load(hf_raw)
-    converted = convert_hf_weights(to_hf_keys(model.state_dict()), hf_cfg)
+    converted = convert_hf_weights(
+        to_hf_keys(model.state_dict(), cfg.hidden_size // cfg.num_attention_heads),
+        hf_cfg,
+    )
     assert_state_dicts_equal(converted, model.state_dict())
 
 
@@ -216,7 +250,9 @@ def test_convert_hf_weights_moe_with_dense_layers_roundtrip():
         decoder_sparse_step=1,
     )
     model = AutoRegressiveLM(cfg)
-    converted = convert_hf_weights(to_hf_keys(model.state_dict()), cfg)
+    converted = convert_hf_weights(
+        to_hf_keys(model.state_dict(), cfg.hidden_size // cfg.num_attention_heads), cfg
+    )
     assert_state_dicts_equal(converted, model.state_dict())
 
 
@@ -230,7 +266,7 @@ def test_convert_hf_weights_qwen2_moe_singular_shared_expert_roundtrip():
         shared_expert_intermediate_size=16,
     )
     model = AutoRegressiveLM(cfg)
-    hf_sd = to_hf_keys(model.state_dict())
+    hf_sd = to_hf_keys(model.state_dict(), cfg.hidden_size // cfg.num_attention_heads)
     hf_sd = {
         k.replace("shared_experts.", "shared_expert.", 1): v for k, v in hf_sd.items()
     }
@@ -241,7 +277,9 @@ def test_convert_hf_weights_qwen2_moe_singular_shared_expert_roundtrip():
 def test_convert_hf_weights_gemma_qk_norm_roundtrip():
     cfg = make_tiny_config(use_qk_norm=True)
     model = AutoRegressiveLM(cfg)
-    converted = convert_hf_weights(to_hf_keys(model.state_dict()), cfg)
+    converted = convert_hf_weights(
+        to_hf_keys(model.state_dict(), cfg.hidden_size // cfg.num_attention_heads), cfg
+    )
     assert_state_dicts_equal(converted, model.state_dict())
 
 
@@ -250,7 +288,9 @@ def test_from_pretrained_hf_directory(tmp_path):
     model = AutoRegressiveLM(cfg).eval()
     save_model(
         config=LLAMA_RAW,
-        state_dict=to_hf_keys(model.state_dict()),
+        state_dict=to_hf_keys(
+            model.state_dict(), cfg.hidden_size // cfg.num_attention_heads
+        ),
         save_directory=str(tmp_path),
     )
     loaded = AutoModel.from_pretrained(tmp_path).eval()
@@ -292,7 +332,9 @@ def test_from_pretrained_weights_format_astrai_rejects_hf(tmp_path):
     model = AutoRegressiveLM(cfg)
     save_model(
         config=LLAMA_RAW,
-        state_dict=to_hf_keys(model.state_dict()),
+        state_dict=to_hf_keys(
+            model.state_dict(), cfg.hidden_size // cfg.num_attention_heads
+        ),
         save_directory=str(tmp_path),
     )
     with pytest.raises(ValueError):
@@ -313,7 +355,7 @@ def test_from_pretrained_invalid_weights_format(tmp_path):
 def test_from_pretrained_hf_directory_sharded(tmp_path):
     cfg = make_tiny_config()
     model = AutoRegressiveLM(cfg).eval()
-    hf_sd = to_hf_keys(model.state_dict())
+    hf_sd = to_hf_keys(model.state_dict(), cfg.hidden_size // cfg.num_attention_heads)
     keys = sorted(hf_sd)
     split = len(keys) // 2
     shard_a = {k: hf_sd[k] for k in keys[:split]}
@@ -342,6 +384,128 @@ def test_from_pretrained_hf_directory_sharded(tmp_path):
         )
 
 
+def _half_split_rope(q, theta=10000.0):
+    """HF llama-style rotate_half RoPE on [batch, seq, heads, head_dim]."""
+    b, s, h, d = q.shape
+    inv_freq = theta ** (-torch.arange(0, d, 2, dtype=torch.float64) / d)
+    freqs = torch.outer(torch.arange(s, dtype=torch.float64), inv_freq).float()
+    cos, sin = freqs.cos()[None, :, None, :], freqs.sin()[None, :, None, :]
+    q1, q2 = q[..., : d // 2], q[..., d // 2 :]
+    return torch.cat([q1 * cos - q2 * sin, q2 * cos + q1 * sin], dim=-1)
+
+
+def _rms_norm_hf(t, weight, eps):
+    t = t.float()
+    t = t * torch.rsqrt(t.pow(2).mean(-1, keepdim=True) + eps)
+    return weight.float() * t
+
+
+def _hf_reference_attn(
+    x, Wq, Wk, Wv, Wo, n_heads, n_kv, head_dim, q_norm_w=None, k_norm_w=None, eps=1e-5
+):
+    """Ground-truth HF attention: per-head RMSNorm BEFORE RoPE (half-split)."""
+    import torch.nn.functional as F
+
+    b, s, dim = x.shape
+    q = (x @ Wq.T).reshape(b, s, n_heads, head_dim).float()
+    k = (x @ Wk.T).reshape(b, s, n_kv, head_dim).float()
+    v = (x @ Wv.T).reshape(b, s, n_kv, head_dim).float()
+    if q_norm_w is not None:
+        q = _rms_norm_hf(q, q_norm_w, eps)
+        k = _rms_norm_hf(k, k_norm_w, eps)
+    q, k = _half_split_rope(q), _half_split_rope(k)
+    rep = n_heads // n_kv
+    k = k.repeat_interleave(rep, dim=2).transpose(1, 2)
+    v = v.repeat_interleave(rep, dim=2).transpose(1, 2)
+    out = F.scaled_dot_product_attention(q.transpose(1, 2), k, v, is_causal=True)
+    out = out.transpose(1, 2).reshape(b, s, n_heads * head_dim)
+    return out @ Wo.T
+
+
+def _run_converted_gqa(x, hf_sd, cfg):
+    from astrai.model.components.attention import GQA
+    from astrai.model.components.rope import get_rotary_emb
+
+    attn = GQA(
+        dim=cfg.hidden_size,
+        n_heads=cfg.num_attention_heads,
+        n_kv_heads=cfg.num_key_value_heads,
+        use_qk_norm=cfg.use_qk_norm,
+        norm_eps=cfg.rms_norm_eps,
+        use_gated_attention=False,
+        layer_id=0,
+    ).eval()
+    converted = convert_hf_weights(hf_sd, cfg)
+    local = {
+        k.removeprefix("layers.0.attention."): v
+        for k, v in converted.items()
+        if k.startswith("layers.0.attention.")
+    }
+    attn.load_state_dict(local, strict=True)
+    head_dim = cfg.hidden_size // cfg.num_attention_heads
+    seq = x.shape[1]
+    rot = get_rotary_emb(head_dim, seq)[None, :seq].expand(x.shape[0], seq, -1, -1)
+    with torch.no_grad():
+        return attn(x, rot, is_causal=True)
+
+
+def test_hf_import_rope_permutation_matches_half_split_reference():
+    torch.manual_seed(0)
+    n_heads, n_kv, head_dim = 4, 2, 8
+    dim = n_heads * head_dim
+    Wq = torch.randn(n_heads * head_dim, dim)
+    Wk = torch.randn(n_kv * head_dim, dim)
+    Wv = torch.randn(n_kv * head_dim, dim)
+    Wo = torch.randn(dim, dim)
+    x = torch.randn(2, 16, dim)
+
+    hf_sd = {
+        "model.layers.0.self_attn.q_proj.weight": Wq,
+        "model.layers.0.self_attn.k_proj.weight": Wk,
+        "model.layers.0.self_attn.v_proj.weight": Wv,
+        "model.layers.0.self_attn.o_proj.weight": Wo,
+    }
+    cfg = make_tiny_config(
+        hidden_size=dim, num_attention_heads=n_heads, num_key_value_heads=n_kv
+    )
+    ref = _hf_reference_attn(x, Wq, Wk, Wv, Wo, n_heads, n_kv, head_dim)
+    out = _run_converted_gqa(x, hf_sd, cfg)
+    torch.testing.assert_close(out, ref, rtol=1e-4, atol=1e-4)
+
+
+def test_hf_import_qk_norm_matches_norm_before_rope_reference():
+    torch.manual_seed(1)
+    n_heads, n_kv, head_dim = 4, 2, 8
+    dim = n_heads * head_dim
+    Wq = torch.randn(n_heads * head_dim, dim)
+    Wk = torch.randn(n_kv * head_dim, dim)
+    Wv = torch.randn(n_kv * head_dim, dim)
+    Wo = torch.randn(dim, dim)
+    gq = torch.randn(head_dim)
+    gk = torch.randn(head_dim)
+    x = torch.randn(2, 16, dim)
+
+    hf_sd = {
+        "model.layers.0.self_attn.q_proj.weight": Wq,
+        "model.layers.0.self_attn.k_proj.weight": Wk,
+        "model.layers.0.self_attn.v_proj.weight": Wv,
+        "model.layers.0.self_attn.o_proj.weight": Wo,
+        "model.layers.0.self_attn.q_norm.weight": gq,
+        "model.layers.0.self_attn.k_norm.weight": gk,
+    }
+    cfg = make_tiny_config(
+        hidden_size=dim,
+        num_attention_heads=n_heads,
+        num_key_value_heads=n_kv,
+        use_qk_norm=True,
+    )
+    ref = _hf_reference_attn(
+        x, Wq, Wk, Wv, Wo, n_heads, n_kv, head_dim, q_norm_w=gq, k_norm_w=gk
+    )
+    out = _run_converted_gqa(x, hf_sd, cfg)
+    torch.testing.assert_close(out, ref, rtol=1e-4, atol=1e-4)
+
+
 def test_from_pretrained_hf_directory_with_moe(tmp_path):
     cfg = make_tiny_config(
         ffn_type="moe",
@@ -354,7 +518,9 @@ def test_from_pretrained_hf_directory_with_moe(tmp_path):
     model = AutoRegressiveLM(cfg).eval()
     save_model(
         config=MOE_RAW,
-        state_dict=to_hf_keys(model.state_dict()),
+        state_dict=to_hf_keys(
+            model.state_dict(), cfg.hidden_size // cfg.num_attention_heads
+        ),
         save_directory=str(tmp_path),
     )
     loaded = AutoModel.from_pretrained(tmp_path).eval()

tests/trainer/test_callbacks.py

@@ -4,7 +4,11 @@ import torch
 
 from astrai.model.components.decoder_block import DecoderBlock
 from astrai.serialization import Checkpoint
-from astrai.trainer.train_callback import GradientCheckpointingCallback, TrainCallback
+from astrai.trainer.train_callback import (
+    GradientCheckpointingCallback,
+    TrainCallback,
+    _copy_tokenizer_files,
+)
 from astrai.trainer.trainer import Trainer
 from tests.helpers import RandomTokenDataset
 
@@ -174,3 +178,38 @@ def test_checkpoint_captures_completed_optimizer_step(
     assert (
         Path(base_test_env["test_dir"]) / "epoch_0_step_1" / "metric.jsonl"
     ).is_file()
+
+
+def test_checkpoint_snapshots_tokenizer_files(
+    base_test_env, train_config_factory, device, tmp_path
+):
+    """Checkpoints copy tokenizer files from param_path so resume works."""
+    param_dir = tmp_path / "model"
+    param_dir.mkdir()
+    (param_dir / "tokenizer.json").write_text("{}")
+    (param_dir / "tokenizer_config.json").write_text("{}")
+
+    train_config = train_config_factory(
+        model_fn=lambda: base_test_env["model"],
+        dataset=RandomTokenDataset(length=2),
+        test_dir=base_test_env["test_dir"],
+        device=device,
+        batch_per_device=2,
+        ckpt_interval=1,
+    )
+
+    Trainer(train_config).train(param_path=str(param_dir))
+
+    ckpt_dir = Path(base_test_env["test_dir"]) / "epoch_0_step_1"
+    assert (ckpt_dir / "tokenizer.json").is_file()
+    assert (ckpt_dir / "tokenizer_config.json").is_file()
+
+    # Resuming with param_path == checkpoint dir must not raise
+    # (samefile guard).
+    _copy_tokenizer_files(str(ckpt_dir), str(ckpt_dir))
+
+
+def test_copy_tokenizer_files_skips_missing_and_none(tmp_path):
+    _copy_tokenizer_files(None, str(tmp_path))
+    _copy_tokenizer_files(str(tmp_path), str(tmp_path / "out"))
+    assert not (tmp_path / "out").exists() or not any((tmp_path / "out").iterdir())

tests/trainer/test_grpo_strategy.py

@@ -162,3 +162,34 @@ def test_grpo_sync_old_model(grpo_strategy):
         if k in old_sd_after
     )
     assert matches
+
+
+def test_grpo_optimizer_step_syncs_old_model(grpo_strategy):
+    """optimizer_step must refresh old_model after each update."""
+    strategy, device = grpo_strategy
+
+    class _SteppedOptimizer:
+        def step(self):
+            with torch.no_grad():
+                for p in strategy.model.parameters():
+                    p.add_(0.05)
+
+    strategy.optimizer_step(_SteppedOptimizer())
+
+    policy_sd = strategy.model.state_dict()
+    old_sd = strategy.old_model.state_dict()
+    assert all(
+        torch.allclose(policy_sd[k], old_sd[k]) for k in policy_sd if k in old_sd
+    )
+
+
+def test_online_grpo_optimizer_step_skips_sync(grpo_strategy):
+    """old_model=None (online) must not attempt a sync."""
+    strategy, device = grpo_strategy
+    strategy.old_model = None
+
+    class _SteppedOptimizer:
+        def step(self):
+            return None
+
+    strategy.optimizer_step(_SteppedOptimizer())

tests/trainer/test_online_strategy.py

@@ -45,6 +45,7 @@ class _RecordingRunner:
         self._fresh = True
         self.policy_version = result.policy_version
         self.weight_updates = []
+        self.eval_calls = 0
 
     def __call__(self, batch):
         self.calls += 1
@@ -52,6 +53,12 @@ class _RecordingRunner:
         self._fresh = False
         return self.result, fresh
 
+    def evaluate(self, batch):
+        # Mirrors RolloutRunner.evaluate: one-off scoring that never
+        # touches the replay cache or freshness state.
+        self.eval_calls += 1
+        return self.result
+
     def step(self):
         self.step_calls += 1
 
@@ -360,6 +367,26 @@ def test_loss_is_differentiable_dpo(device):
     assert has_grad
 
 
+def test_validate_online_returns_none_without_runner(device):
+    strat = _make_grpo(device)
+    batch = {"input_ids": torch.randint(3, 200, (2, 4), device=device)}
+    assert strat.validate_online(batch) is None
+
+
+def test_validate_online_uses_one_off_rollout_not_replay_cache(device):
+    strat = _make_grpo(device)
+    runner = _RecordingRunner(_make_rollout_result(device=device))
+    strat.set_rollout_runner(runner)
+
+    out = strat.validate_online(
+        {"input_ids": torch.randint(3, 200, (2, 4), device=device)}
+    )
+
+    assert torch.isfinite(out["loss"]).item()
+    assert runner.eval_calls == 1
+    assert runner.calls == 0  # replay cache path untouched
+
+
 def test_ref_model_not_updated_by_backward_dpo(device):
     strat = _make_dpo(device)
     strat.set_rollout_runner(_RecordingRunner(_make_rollout_result(device=device)))

tests/trainer/test_rollout.py

@@ -388,6 +388,21 @@ def test_rollout_runner_cache_returns_stale_flag(device):
     assert fresh2 is False
 
 
+def test_rollout_runner_evaluate_leaves_cache_untouched(device):
+    runner, _ = _make_runner(device, rollout_interval=10)
+    batch = _make_instruction_batch()
+    cached, _ = runner(batch)
+
+    eval_batch = _make_instruction_batch(n=1)
+    result = runner.evaluate(eval_batch)
+
+    assert result.rewards.shape == result.responses.shape[:2]
+    replayed, fresh = runner(batch)
+    assert replayed is cached
+    assert fresh is False
+    assert runner._steps_since_rollout == 0
+
+
 def test_rollout_runner_tags_generation_version_and_preserves_cached_behavior(device):
     runner, _ = _make_runner(device, rollout_interval=100)
     batch = _make_instruction_batch(n=1)