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_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())