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

@@ -17,6 +17,7 @@ from astrai.parallel.setup import (
     setup_parallel,
     spawn_parallel_fn,
 )
+from astrai.parallel.utils import create_ref_model
 
 __all__ = [
     "get_world_size",
@@ -35,4 +36,5 @@ __all__ = [
     "NoneExecutor",
     "DDPExecutor",
     "FSDPExecutor",
+    "create_ref_model",
 ]

astrai/parallel/executor.py

@@ -326,21 +326,27 @@ class FSDPExecutor(BaseExecutor):
         if not self.use_distributed:
             return model.state_dict()
 
-        if get_rank() != 0:
-            return None
-
+        # unshard() and full_tensor() are collective ops — all ranks must
+        # participate. Non-rank-0 ranks still call them but discard results.
         for module in model.modules():
             if isinstance(module, FSDPModule):
                 module.unshard()
 
         state_dict = model.state_dict()
-        result = {
-            k: (v.full_tensor() if isinstance(v, DTensor) else v)
-            for k, v in state_dict.items()
-        }
+        result = {}
+        for k, v in state_dict.items():
+            if isinstance(v, DTensor):
+                full = v.full_tensor()
+                if get_rank() == 0:
+                    result[k] = full
+            elif get_rank() == 0:
+                result[k] = v
 
         for module in model.modules():
             if isinstance(module, FSDPModule):
                 module.reshard()
 
+        if get_rank() != 0:
+            return None
+
         return result

astrai/parallel/utils.py

@@ -0,0 +1,34 @@
+"""Utility functions for parallel training."""
+
+from typing import TYPE_CHECKING, Callable, Dict, Optional
+
+import torch
+import torch.nn as nn
+
+if TYPE_CHECKING:
+    from astrai.parallel.executor import BaseExecutor
+
+
+def create_ref_model(
+    model_fn: Callable[[], nn.Module],
+    executor: Optional["BaseExecutor"] = None,
+    model: Optional[nn.Module] = None,
+    state_dict: Optional[Dict[str, torch.Tensor]] = None,
+    device: Optional[str] = None,
+) -> Optional[nn.Module]:
+    """Create a frozen reference model from executor or state dict.
+
+    On non-rank-0, returns None (executor.unwrap_model returns None).
+    """
+    if state_dict is None and executor is not None and model is not None:
+        state_dict = executor.unwrap_model(model)
+    if state_dict is None:
+        return None
+
+    ref_model = model_fn()
+    ref_model.load_state_dict(state_dict)
+    ref_model.requires_grad_(False)
+    ref_model.eval()
+    if device is not None:
+        ref_model = ref_model.to(device=device)
+    return ref_model

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)

astrai/trainer/train_context.py

@@ -14,11 +14,12 @@ from astrai.inference.core.scheduler import InferenceScheduler
 from astrai.model.components.lora import inject_lora
 from astrai.parallel.executor import BaseExecutor, ExecutorFactory
 from astrai.parallel.setup import get_current_device, get_rank, get_world_size
+from astrai.parallel.utils import create_ref_model
 from astrai.protocols import OptimizerProtocol, SchedulerProtocol
 from astrai.serialization import Checkpoint, load_json
 from astrai.tokenize import AutoTokenizer
 from astrai.trainer.rollout import RolloutGenerator, RolloutRunner
-from astrai.trainer.strategy import BaseStrategy, StrategyFactory, create_ref_model
+from astrai.trainer.strategy import BaseStrategy, StrategyFactory
 
 logger = logging.getLogger(__name__)
 
@@ -229,17 +230,14 @@ class TrainContextBuilder:
         needs_old = cfg.strategy in ("grpo", "online_grpo")
 
         if needs_ref:
-            ref_model = create_ref_model(
-                cfg.model_fn, executor.unwrap_model(context.model)
-            ).to(device=device)
-            strategy_kwargs["ref_model"] = ref_model
+            strategy_kwargs["ref_model"] = create_ref_model(
+                cfg.model_fn, executor=executor, model=context.model, device=device
+            )
 
-        old_model = None
         if needs_old:
-            old_model = create_ref_model(
-                cfg.model_fn, executor.unwrap_model(context.model)
-            ).to(device=device)
-            strategy_kwargs["old_model"] = old_model
+            strategy_kwargs["old_model"] = create_ref_model(
+                cfg.model_fn, executor=executor, model=context.model, device=device
+            )
 
         context.strategy = StrategyFactory.create(
             cfg.strategy,