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/parallel/executor.py

@@ -61,12 +61,13 @@ def broadcast_state_dict(
 
     rank = dist.get_rank()
 
-    # Broadcast metadata (keys, shapes, dtypes, device) so non-src ranks
-    # can allocate matching empty tensors on the correct device.
+    # Broadcast metadata (keys, shapes, dtypes, device kind) so non-src
+    # ranks can allocate matching empty tensors on their own device.
     if rank == src:
-        device = next(iter(state_dict.values())).device
+        first = next(iter(state_dict.values()), None)
+        on_cuda = first is not None and first.is_cuda
         metadata = [
-            (k, tuple(v.shape), v.dtype, str(device)) for k, v in state_dict.items()
+            (k, tuple(v.shape), v.dtype, on_cuda) for k, v in state_dict.items()
         ]
     else:
         metadata = None
@@ -75,10 +76,18 @@ def broadcast_state_dict(
     metadata = metadata_list[0]
 
     # Non-src ranks allocate empty tensors with the broadcasted metadata.
+    # The allocation must sit on *this* rank's local device: process groups
+    # are pinned per-rank via ``device_id=``, so NCCL rejects a tensor
+    # allocated on the src rank's device (and the cross-GPU allocation
+    # would be wrong here anyway).
     if rank != src:
+        device = (
+            torch.device("cuda", torch.cuda.current_device())
+            if any(m[3] for m in metadata)
+            else torch.device("cpu")
+        )
         state_dict = {
-            k: torch.empty(s, dtype=d, device=torch.device(dev))
-            for k, s, d, dev in metadata
+            k: torch.empty(s, dtype=d, device=device) for k, s, d, _ in metadata
         }
 
     # Broadcast each tensor in-place.
@@ -400,17 +409,16 @@ class FSDPExecutor(BaseExecutor):
         if not self.use_distributed:
             return super().clip_grad_norm(model, max_norm)
 
-        # FSDP params are DTensors (sharded across ranks).
-        # torch.nn.utils.clip_grad_norm_ computes LOCAL norm per rank,
-        # so we must all-reduce to get the global norm before clipping.
-        local_norm = torch.nn.utils.get_total_norm(
+        # FSDP params are DTensors (sharded across ranks), and
+        # get_total_norm reduces them to a globally-reduced replicated
+        # DTensor — no extra all-reduce is needed. Manually summing the
+        # local slice again would inflate the norm by sqrt(world_size)
+        # and over-clip every step.
+        total_norm = torch.nn.utils.get_total_norm(
             [p.grad for p in model.parameters() if p.grad is not None],
         )
-        if isinstance(local_norm, DTensor):
-            local_norm = local_norm.to_local()
-        total_norm_sq = local_norm**2
-        dist.all_reduce(total_norm_sq, op=dist.ReduceOp.SUM)
-        total_norm = total_norm_sq.sqrt()
+        if isinstance(total_norm, DTensor):
+            total_norm = total_norm.full_tensor()
 
         clip_coef = max_norm / (total_norm + 1e-6)
         clip_coef_clamped = torch.clamp(clip_coef, max=1.0)