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/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"]