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

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)