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/model/components/attention.py

@@ -69,11 +69,15 @@ class GQA(nn.Module):
         q = self._split_heads(self.q_proj(x), self.n_heads)
         k = self._split_heads(self.k_proj(x), self.n_kv_heads)
         v = self._split_heads(self.v_proj(x), self.n_kv_heads)
-        q, k = apply_rotary_emb(q, rotary_emb), apply_rotary_emb(k, rotary_emb)
 
+        # Match the HuggingFace convention (Qwen2/Gemma): normalize Q/K
+        # before RoPE. RMSNorm's per-channel gain does not commute with
+        # the rotation, so the order changes numerics.
         if self.use_qk_norm:
             q, k = self.q_norm(q), self.k_norm(k)
 
+        q, k = apply_rotary_emb(q, rotary_emb), apply_rotary_emb(k, rotary_emb)
+
         sdqa_out = attention(
             q, k, v, kv_cache, self.layer_id, attn_mask, is_causal, fwd
         ).reshape(*x.shape[:-1], self.dim)