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