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",
 ]