diff --git a/astrai/parallel/__init__.py b/astrai/parallel/__init__.py index b15f85c..b565f8d 100644 --- a/astrai/parallel/__init__.py +++ b/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", ] diff --git a/astrai/parallel/executor.py b/astrai/parallel/executor.py index b8523e5..d0c12c8 100644 --- a/astrai/parallel/executor.py +++ b/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 diff --git a/astrai/trainer/strategy.py b/astrai/trainer/strategy.py index 6f7bb2f..a205061 100644 --- a/astrai/trainer/strategy.py +++ b/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) diff --git a/tests/conftest.py b/tests/conftest.py index be6abaa..47f302d 100644 --- a/tests/conftest.py +++ b/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") diff --git a/tests/extension/conftest.py b/tests/extension/conftest.py index 4af48a8..5f398fd 100644 --- a/tests/extension/conftest.py +++ b/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 - -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" -) +from tests.conftest import skip_no_kernel D = 64 CFG = dict( diff --git a/tests/extension/test_backend_equivalence.py b/tests/extension/test_backend_equivalence.py index bc88ea1..0d83520 100644 --- a/tests/extension/test_backend_equivalence.py +++ b/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 diff --git a/tests/extension/test_kernel_mask.py b/tests/extension/test_kernel_mask.py index 078bd21..22d9195 100644 --- a/tests/extension/test_kernel_mask.py +++ b/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 diff --git a/tests/helpers.py b/tests/helpers.py index 17c4428..474977e 100644 --- a/tests/helpers.py +++ b/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 diff --git a/tests/parallel/test_broadcast_state_dict.py b/tests/parallel/test_broadcast_state_dict.py new file mode 100644 index 0000000..4d5a349 --- /dev/null +++ b/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])