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/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(