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/preprocessing/pipeline.py

@@ -146,11 +146,17 @@ class Pipeline:
 
     @staticmethod
     def _align_bucket(bucket: dict, result: dict, ids: list):
-        """Pad previously-accumulated keys that are missing from *result*."""
+        """Pad previously-accumulated keys that are missing from *result*.
+
+        Builders omit all-ones masks (``loss_mask`` / ``*_mask``) to save
+        space, so an omitted mask means "train on every token" and is
+        back-filled with ones; every other missing key pads with zeros.
+        """
         for key in list(bucket.keys()):
             if key in result:
                 continue
-            bucket[key].append([0] * len(ids))
+            fill = 1 if key == "loss_mask" or key.endswith("_mask") else 0
+            bucket[key].append([fill] * len(ids))
 
     def _iter_items(self):
         for path in self.paths: