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