fix: FSDP unwrap_model collective op and None guard

- unshard() and full_tensor() are collective ops, all ranks must participate
- Old code returned None on non-rank-0 before calling unshard, causing deadlock
- Fix: all ranks unshard/full_tensor, only rank-0 keeps the result
- Move create_ref_model to parallel/utils.py, accept executor+model directly
- Guard create_ref_model and sync_old_model against None on non-rank-0

astrai/trainer/strategy.py

@@ -9,20 +9,10 @@ import torch.nn.functional as F
 from torch import Tensor
 
 from astrai.factory import BaseFactory
+from astrai.parallel.utils import create_ref_model
 from astrai.trainer.rollout import RolloutResult
 
 
-def create_ref_model(
-    model_fn: Callable[[], nn.Module], state_dict: Dict[str, Tensor]
-) -> nn.Module:
-    """Create a frozen reference model from model_fn + full state dict."""
-    ref_model = model_fn()
-    ref_model.load_state_dict(state_dict)
-    ref_model.requires_grad_(False)
-    ref_model.eval()
-    return ref_model
-
-
 def move_to_device(batch: Dict[str, Tensor], device: str) -> Dict[str, Tensor]:
     """Move batch tensors to specified device with non-blocking transfer."""
     return {key: value.to(device, non_blocking=True) for key, value in batch.items()}
@@ -401,7 +391,9 @@ class GRPOStrategy(BaseStrategy):
 
     def sync_old_model(self):
         """Copy current policy weights to old model."""
-        self.old_model.load_state_dict(self.executor.unwrap_model(self.model))
+        state_dict = self.executor.unwrap_model(self.model)
+        if state_dict is not None:
+            self.old_model.load_state_dict(state_dict)
 
     def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
         batch = move_to_device(batch, self.device)