fix: broadcast ref/old model state_dict for FSDP

- Add broadcast_state_dict to sync state_dict from rank-0 to all ranks
- Fix create_ref_model returning None on non-rank-0 under FSDP
- Fix sync_old_model only updating old_model on rank-0 under FSDP
- Split skip_no_cuda/skip_no_kernel markers and hoist to top-level conftest
- Add distributed tests for broadcast_state_dict and create_ref_model

astrai/parallel/__init__.py

@@ -7,6 +7,7 @@ from astrai.parallel.executor import (
     FSDPExecutor,
     GradientState,
     NoneExecutor,
+    broadcast_state_dict,
     create_ref_model,
 )
 from astrai.parallel.setup import (
@@ -34,4 +35,5 @@ __all__ = [
     "DDPExecutor",
     "FSDPExecutor",
     "create_ref_model",
+    "broadcast_state_dict",
 ]

astrai/parallel/executor.py

@@ -24,6 +24,49 @@ from astrai.parallel.setup import get_rank, get_world_size
 logger = logging.getLogger(__name__)
 
 
+def broadcast_state_dict(
+    state_dict: Optional[Dict[str, torch.Tensor]],
+    src: int = 0,
+) -> Optional[Dict[str, torch.Tensor]]:
+    """Broadcast a state_dict from *src* rank to all ranks.
+
+    Tensors stay on their original device (GPU) for the broadcast.
+    All ranks must call this collectively.
+
+    On non-distributed runs, returns *state_dict* unchanged.
+    """
+    if not dist.is_initialized() or dist.get_world_size() == 1:
+        return 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.
+    if rank == src:
+        device = next(iter(state_dict.values())).device
+        metadata = [
+            (k, tuple(v.shape), v.dtype, str(device)) for k, v in state_dict.items()
+        ]
+    else:
+        metadata = None
+    metadata_list = [metadata]
+    dist.broadcast_object_list(metadata_list, src=src)
+    metadata = metadata_list[0]
+
+    # Non-src ranks allocate empty tensors with the broadcasted metadata.
+    if rank != src:
+        state_dict = {
+            k: torch.empty(s, dtype=d, device=torch.device(dev))
+            for k, s, d, dev in metadata
+        }
+
+    # Broadcast each tensor in-place.
+    for tensor in state_dict.values():
+        dist.broadcast(tensor, src=src)
+
+    return state_dict
+
+
 def create_ref_model(
     model_fn: Callable[[], nn.Module],
     executor: Optional["BaseExecutor"] = None,
@@ -33,10 +76,18 @@ def create_ref_model(
 ) -> Optional[nn.Module]:
     """Create a frozen reference model from executor or state dict.
 
-    On non-rank-0, returns None (executor.unwrap_model returns None).
+    In distributed mode (FSDP), ``unwrap_model`` returns ``None`` on
+    non-rank-0.  The state_dict is broadcast from rank-0 to all ranks
+    so every rank gets a complete copy.
     """
     if state_dict is None and executor is not None and model is not None:
         state_dict = executor.unwrap_model(model)
+
+    # FSDP's unwrap_model returns None on non-rank-0. Broadcast from
+    # rank-0 so every rank receives a complete state_dict.
+    if executor is not None and executor.use_distributed:
+        state_dict = broadcast_state_dict(state_dict)
+
     if state_dict is None:
         return None
 

astrai/trainer/strategy.py

@@ -9,6 +9,7 @@ import torch.nn.functional as F
 from torch import Tensor
 
 from astrai.factory import BaseFactory
+from astrai.parallel.executor import broadcast_state_dict
 from astrai.trainer.rollout import RolloutResult
 
 
@@ -391,6 +392,8 @@ class GRPOStrategy(BaseStrategy):
     def sync_old_model(self):
         """Copy current policy weights to old model."""
         state_dict = self.executor.unwrap_model(self.model)
+        if self.executor.use_distributed:
+            state_dict = broadcast_state_dict(state_dict)
         if state_dict is not None:
             self.old_model.load_state_dict(state_dict)
 

tests/conftest.py

@@ -7,10 +7,16 @@ import pytest
 import torch
 from tokenizers import Tokenizer, models, pre_tokenizers, trainers
 
+from astrai.extension import KERNEL_NAMES, is_available
 from astrai.model.transformer import AutoRegressiveLM
 from astrai.tokenize import AutoTokenizer
 from tests.helpers import TINY_CONFIG, RandomTokenDataset, make_tiny_config
 
+CUDA_AVAIL = torch.cuda.is_available()
+KERNEL_AVAIL = CUDA_AVAIL and all(is_available(k) for k in KERNEL_NAMES)
+skip_no_cuda = pytest.mark.skipif(not CUDA_AVAIL, reason="CUDA not available")
+skip_no_kernel = pytest.mark.skipif(not KERNEL_AVAIL, reason="CUDA kernels not built")
+
 
 def pytest_configure(config):
     config.addinivalue_line("markers", "slow: marks tests as slow")

tests/extension/conftest.py

@@ -4,13 +4,8 @@ import pytest
 import torch
 
 from astrai.config.model_config import AutoRegressiveLMConfig
-from astrai.extension import is_available
 from astrai.model.transformer import AutoRegressiveLM
-
+from tests.conftest import skip_no_kernel
-CUDA_AVAILABLE = torch.cuda.is_available() and is_available("attn_paged_decode")
-skip_no_cuda = pytest.mark.skipif(
-    not CUDA_AVAILABLE, reason="CUDA not available or kernels not built"
-)
 
 D = 64
 CFG = dict(

tests/extension/test_backend_equivalence.py

@@ -8,10 +8,10 @@ import torch
 
 from astrai.extension import ATTN_BACKEND, attn_backend
 from astrai.inference.core.cache import PagePool
-from tests.extension.conftest import D, skip_no_cuda
+from tests.extension.conftest import D, skip_no_kernel
 
 
-@skip_no_cuda
+@skip_no_kernel
 def test_training_forward_matches_torch(cuda_model):
     """Training forward (kv_cache=None) should produce identical logits."""
     model, _ = cuda_model
@@ -27,7 +27,7 @@ def test_training_forward_matches_torch(cuda_model):
     assert diff == 0.0, f"Training forward diff {diff} should be 0"
 
 
-@skip_no_cuda
+@skip_no_kernel
 def test_prefill_with_kv_cache_matches_torch(cuda_model):
     """Inference prefill with KV cache should match torch backend."""
     model, _ = cuda_model
@@ -93,7 +93,7 @@ def test_prefill_with_kv_cache_matches_torch(cuda_model):
         assert d == 0.0, f"Prefill diff for sample {i}: {d}"
 
 
-@skip_no_cuda
+@skip_no_kernel
 def test_decode_mixed_seq_lens_matches_torch(cuda_model):
     """Decode with mixed seq_lens in batch — padding mask must produce correct output."""
     model, _ = cuda_model
@@ -152,7 +152,7 @@ def test_decode_mixed_seq_lens_matches_torch(cuda_model):
     assert diff < 0.05, f"Decode diff (mixed seq_lens): {diff}"
 
 
-@skip_no_cuda
+@skip_no_kernel
 def test_run_batch_cuda_matches_torch_greedy(cuda_model):
     """Greedy decode (temperature=0) should produce identical tokens."""
     from astrai.inference.core.scheduler import InferenceScheduler

tests/extension/test_kernel_mask.py

@@ -2,10 +2,10 @@
 
 import torch
 
-from tests.extension.conftest import D, skip_no_cuda
+from tests.extension.conftest import D, skip_no_kernel
 
 
-@skip_no_cuda
+@skip_no_kernel
 def test_kernel_accepts_2d_mask():
     """Kernel should accept 2D mask [batch, kv_len]."""
     from astrai.extension.attention_ops import attn_prefill
@@ -22,7 +22,7 @@ def test_kernel_accepts_2d_mask():
     assert out.shape == (batch, q_len, n_heads, D)
 
 
-@skip_no_cuda
+@skip_no_kernel
 def test_kernel_accepts_3d_mask():
     """Kernel should accept 3D mask [batch, q_len, kv_len]."""
     from astrai.extension.attention_ops import attn_prefill
@@ -38,7 +38,7 @@ def test_kernel_accepts_3d_mask():
     assert out.shape == (batch, q_len, n_heads, D)
 
 
-@skip_no_cuda
+@skip_no_kernel
 def test_kernel_accepts_4d_mask():
     """Kernel should accept 4D mask [batch, n_heads, q_len, kv_len]."""
     from astrai.extension.attention_ops import attn_prefill
@@ -55,7 +55,7 @@ def test_kernel_accepts_4d_mask():
     assert out.shape == (batch, q_len, n_heads, D)
 
 
-@skip_no_cuda
+@skip_no_kernel
 def test_4d_mask_matches_no_mask_when_all_true():
     """A 4D all-True mask should produce the same output as no mask."""
     from astrai.extension.attention_ops import attn_prefill

tests/helpers.py

@@ -165,6 +165,8 @@ class FakeTokenizer:
 class FakeExecutor:
     """Executor stub tracking ``sync_gradients`` and providing ``unwrap_model``."""
 
+    use_distributed = False
+
     def __init__(self, sync_gradients=True):
         self._sync_gradients = sync_gradients
 

tests/parallel/test_broadcast_state_dict.py

@@ -0,0 +1,172 @@
+"""Tests for :func:`broadcast_state_dict` and distributed ``create_ref_model``.
+
+Uses ``spawn_parallel_fn`` with the ``gloo`` backend to simulate a
+multi-rank environment without requiring multiple GPUs.
+"""
+
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+
+from astrai.model.transformer import AutoRegressiveLM
+from astrai.parallel import get_rank, spawn_parallel_fn
+from astrai.parallel.executor import broadcast_state_dict, create_ref_model
+from astrai.trainer.strategy import GRPOStrategy
+from tests.helpers import FakeExecutor, make_rollout_config
+
+
+def _broadcast_worker():
+    """Rank-0 builds a state_dict; all ranks verify they receive it."""
+    rank = get_rank()
+
+    if rank == 0:
+        sd = {
+            "layer.weight": torch.randn(4, 8),
+            "layer.bias": torch.randn(4),
+        }
+        expected = {k: v.clone() for k, v in sd.items()}
+    else:
+        sd = None
+        expected = None
+
+    received = broadcast_state_dict(sd, src=0)
+
+    assert received is not None, f"rank {rank}: received None"
+    assert set(received.keys()) == {"layer.weight", "layer.bias"}
+    if rank == 0:
+        # rank-0 already had the data
+        for k in received:
+            assert torch.equal(received[k], expected[k])
+    # tensors preserve the source device (cpu here since gloo test)
+    for k, v in received.items():
+        assert v.device.type == "cpu", f"rank {rank}: {k} on {v.device}"
+
+
+def test_broadcast_state_dict():
+    spawn_parallel_fn(_broadcast_worker, world_size=2, backend="gloo")
+
+
+def _create_ref_model_worker():
+    """Verify create_ref_model works when unwrap_model returns None on non-rank-0."""
+
+    class FakeFSDPExecutor:
+        """Simulates FSDP: unwrap_model returns state_dict on rank-0, None elsewhere."""
+
+        use_distributed = True
+
+        def unwrap_model(self, model):
+            if get_rank() == 0:
+                return model.state_dict()
+            return None
+
+    rank = get_rank()
+    config = make_rollout_config()
+    model = AutoRegressiveLM(config).to("cpu")
+    # Give each rank distinct weights so we can verify broadcast overwrites them
+    with torch.no_grad():
+        for p in model.parameters():
+            p.add_(float(rank))
+
+    executor = FakeFSDPExecutor()
+
+    ref = create_ref_model(
+        model_fn=lambda: AutoRegressiveLM(config),
+        executor=executor,
+        model=model,
+        device="cpu",
+    )
+
+    assert ref is not None, f"rank {rank}: ref model is None"
+    # Every rank should have rank-0's weights, not its own
+    rank0_sd = model.state_dict() if rank == 0 else None
+    # Broadcast rank-0's original weights for comparison
+    if rank == 0:
+        expected_sd = {k: v.clone() for k, v in model.state_dict().items()}
+    else:
+        expected_sd = None
+    expected_sd = broadcast_state_dict(expected_sd, src=0)
+
+    ref_sd = ref.state_dict()
+    for k in ref_sd:
+        assert torch.equal(ref_sd[k], expected_sd[k]), f"rank {rank}: mismatch at {k}"
+    # ref model should be frozen and in eval mode
+    assert not ref.training
+    for p in ref.parameters():
+        assert not p.requires_grad
+
+
+def test_create_ref_model_distributed():
+    spawn_parallel_fn(_create_ref_model_worker, world_size=2, backend="gloo")
+
+
+def _sync_old_model_worker():
+    """Verify that sync_old_model broadcasts weights to all ranks."""
+    rank = get_rank()
+    config = make_rollout_config()
+    model = AutoRegressiveLM(config).to("cpu")
+    old_model = AutoRegressiveLM(config).to("cpu")
+    ref_model = AutoRegressiveLM(config).to("cpu")
+
+    # Give model rank-distinct weights
+    with torch.no_grad():
+        for p in model.parameters():
+            p.add_(float(rank) * 10)
+
+    class _DistExecutor:
+        use_distributed = True
+
+        def unwrap_model(self, m):
+            if get_rank() == 0:
+                return m.state_dict()
+            return None
+
+    strategy = GRPOStrategy(
+        model=model,
+        device="cpu",
+        old_model=old_model,
+        ref_model=ref_model,
+        executor=_DistExecutor(),
+    )
+
+    # Capture rank-0's policy weights for comparison
+    if rank == 0:
+        expected = {k: v.clone() for k, v in model.state_dict().items()}
+    else:
+        expected = None
+    expected = broadcast_state_dict(expected, src=0)
+
+    strategy.sync_old_model()
+
+    old_sd = strategy.old_model.state_dict()
+    for k in old_sd:
+        assert torch.equal(old_sd[k], expected[k]), f"rank {rank}: mismatch at {k}"
+
+
+def test_sync_old_model_distributed():
+    spawn_parallel_fn(_sync_old_model_worker, world_size=2, backend="gloo")
+
+
+def test_broadcast_state_dict_single_process():
+    """When dist is not initialized, broadcast_state_dict is a no-op."""
+    sd = {"w": torch.randn(3, 3), "b": torch.randn(3)}
+    result = broadcast_state_dict(sd)
+    assert result is sd
+
+
+def test_create_ref_model_single_process():
+    """create_ref_model still works without an executor (explicit state_dict)."""
+    config = make_rollout_config()
+    model = AutoRegressiveLM(config)
+    sd = model.state_dict()
+
+    ref = create_ref_model(
+        model_fn=lambda: AutoRegressiveLM(config),
+        state_dict=sd,
+        device="cpu",
+    )
+    assert ref is not None
+    assert not ref.training
+    for p in ref.parameters():
+        assert not p.requires_grad
+    for k in sd:
+        assert torch.equal(ref.state_dict()[k], sd[k])