fix: support DDP online rollout

astrai/parallel/__init__.py

@@ -7,6 +7,7 @@
     FSDPExecutor,
     GradientState,
     NoneExecutor,
+    RolloutCapabilities,
     broadcast_state_dict,
     create_ref_model,
 )
@@ -34,6 +35,7 @@
     "NoneExecutor",
     "DDPExecutor",
     "FSDPExecutor",
+    "RolloutCapabilities",
     "create_ref_model",
     "broadcast_state_dict",
 ]

astrai/parallel/executor.py

@@ -4,6 +4,7 @@
 import logging
 import os
 from contextlib import contextmanager
+from dataclasses import dataclass
 from typing import Any, Callable, Dict, Optional, Tuple
 
 import torch
@@ -158,9 +159,20 @@ def get_last_lr(self):
         return self.scheduler.get_last_lr()
 
 
+@dataclass(frozen=True)
+class RolloutCapabilities:
+    """Executor capabilities required by online rollout."""
+
+    supported: bool
+    supports_compile: bool = False
+    reason: Optional[str] = None
+
+
 class BaseExecutor:
     def __init__(self, grad_accum_steps: int = 1):
         self.gradient_state = GradientState(grad_accum_steps)
+        self._training_model: Optional[nn.Module] = None
+        self._inference_model: Optional[nn.Module] = None
 
     def prepare(
         self,
@@ -174,8 +186,11 @@ def prepare(
         if before_wrap is not None:
             model = before_wrap(model)
         model = self._prepare_model(model)
+        inference_model = self._prepare_inference_model(model)
         if after_wrap is not None:
             model = after_wrap(model)
+        self._training_model = model
+        self._inference_model = inference_model
         optimizer = None
         scheduler = None
         if optimizer_fn is not None:
@@ -190,6 +205,30 @@ def prepare(
     def _prepare_model(self, model: nn.Module) -> nn.Module:
         return model
 
+    def _prepare_inference_model(self, model: nn.Module) -> nn.Module:
+        return model
+
+    def rollout_capabilities(self) -> RolloutCapabilities:
+        return RolloutCapabilities(supported=True)
+
+    def model_for_training(self) -> nn.Module:
+        """Return the executor-owned model used for loss/backward."""
+        if self._training_model is None:
+            raise RuntimeError("Executor model has not been prepared")
+        return self._training_model
+
+    def model_for_inference(self) -> nn.Module:
+        """Return the supported model view used by online rollout."""
+        capabilities = self.rollout_capabilities()
+        if not capabilities.supported:
+            detail = f": {capabilities.reason}" if capabilities.reason else ""
+            raise RuntimeError(
+                f"{type(self).__name__} does not support online rollout{detail}"
+            )
+        if self._inference_model is None:
+            raise RuntimeError("Executor model has not been prepared")
+        return self._inference_model
+
     def _no_sync(self, model: nn.Module):
         return contextlib.nullcontext()
 
@@ -292,15 +331,27 @@ def _prepare_model(self, model: nn.Module) -> nn.Module:
             logger.warning("DDP backend selected but world_size=1, model not wrapped")
             return model
         local_rank = int(os.environ.get("LOCAL_RANK", get_rank()))
-        model = DDP(
-            model,
-            device_ids=[local_rank],
-            output_device=local_rank,
-            **self._ddp_kwargs,
-        )
+        try:
+            device = next(model.parameters()).device
+        except StopIteration:
+            device = torch.device("cpu")
+        if device.type == "cpu":
+            model = DDP(model, **self._ddp_kwargs)
+        else:
+            model = DDP(
+                model,
+                device_ids=[local_rank],
+                output_device=local_rank,
+                **self._ddp_kwargs,
+            )
         logger.info("Model wrapped with DDP (world_size=%d)", get_world_size())
         return model
 
+    def _prepare_inference_model(self, model: nn.Module) -> nn.Module:
+        if isinstance(model, DDP):
+            return model.module
+        return model
+
     def _no_sync(self, model: nn.Module):
         if isinstance(model, DDP):
             return model.no_sync()
@@ -335,6 +386,15 @@ def __init__(
         self._mp_policy = mp_policy
         self._reshard_after_forward = reshard_after_forward
 
+    def rollout_capabilities(self) -> RolloutCapabilities:
+        return RolloutCapabilities(
+            supported=False,
+            reason=(
+                "FSDP-sharded parameters are not supported by the local "
+                "InferenceScheduler"
+            ),
+        )
+
     def _prepare_model(self, model: nn.Module) -> nn.Module:
         if not self.use_distributed:
             logger.warning("FSDP backend selected but world_size=1, model not wrapped")

astrai/trainer/train_context.py

@@ -97,11 +97,13 @@ def with_param_path(self, param_path: Optional[str], resume: bool = False) -> Se
         return self
 
     def build(self) -> TrainContext:
+        executor = self._create_executor()
+        self._validate_rollout_execution(executor)
+
         # Resolve persisted state.
         preloaded_state = self._load_preloaded_state()
 
         # Build the core training components and restore their persisted state.
-        executor = self._create_executor()
         context = self._create_context(preloaded_state, executor)
         self._prepare_model(context, executor, preloaded_state)
         self._restore_optimizer_state(context)
@@ -124,6 +126,24 @@ def _create_executor(self) -> BaseExecutor:
             **cfg.executor_kwargs,
         )
 
+    def _validate_rollout_execution(self, executor: BaseExecutor) -> None:
+        cfg = self.config
+        if not cfg.strategy.startswith("online_"):
+            return
+
+        capabilities = executor.rollout_capabilities()
+        if not capabilities.supported:
+            detail = f": {capabilities.reason}" if capabilities.reason else ""
+            raise ValueError(
+                "Online rollout is not supported with "
+                f"parallel_mode='{cfg.parallel_mode}'{detail}"
+            )
+        if cfg.compile_mode is not None and not capabilities.supports_compile:
+            raise ValueError(
+                "Online rollout with torch.compile is not supported with "
+                f"parallel_mode='{cfg.parallel_mode}'"
+            )
+
     def _load_preloaded_state(self) -> _PreloadedState:
         cfg = self.config
         state = _PreloadedState(
@@ -312,12 +332,13 @@ def _configure_rollout(self, context: TrainContext, strategy_kwargs: dict) -> No
                 f"Strategy '{cfg.strategy}' does not support online rollout"
             )
         tokenizer = AutoTokenizer.from_pretrained(self._param_path)
+        rollout_model = context.executor.model_for_inference()
         group_size = strategy_kwargs.get("group_size", 1)
         scheduler = InferenceScheduler(
-            model=context.model,
+            model=rollout_model,
             tokenizer=tokenizer,
             max_batch_size=group_size * max(1, cfg.batch_per_device),
-            max_seq_len=getattr(context.model.config, "max_position_embeddings", None),
+            max_seq_len=getattr(rollout_model.config, "max_position_embeddings", None),
         )
         generator = RolloutGenerator(
             scheduler=scheduler,

tests/parallel/test_parallel.py

@@ -1,7 +1,16 @@
+import pytest
 import torch
 import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel as DDP
 
-from astrai.parallel import get_rank, only_on_rank, spawn_parallel_fn
+from astrai.parallel import (
+    DDPExecutor,
+    FSDPExecutor,
+    NoneExecutor,
+    get_rank,
+    only_on_rank,
+    spawn_parallel_fn,
+)
 
 
 @only_on_rank(0)
@@ -24,9 +33,48 @@ def all_reduce():
     assert x.item() == expected_sum
 
 
+def ddp_model_views():
+    executor = DDPExecutor()
+    training_model, _, _ = executor.prepare(lambda: torch.nn.Linear(4, 3))
+
+    assert isinstance(training_model, DDP)
+    assert executor.model_for_training() is training_model
+    assert executor.model_for_inference() is training_model.module
+    assert executor.rollout_capabilities().supported
+
+    inputs = torch.arange(8, dtype=torch.float32).reshape(2, 4)
+    training_choice = training_model(inputs).argmax(dim=-1)
+    inference_choice = executor.model_for_inference()(inputs).argmax(dim=-1)
+    assert torch.equal(training_choice, inference_choice)
+
+
+def test_none_executor_exposes_same_training_and_inference_model():
+    executor = NoneExecutor()
+    model, _, _ = executor.prepare(lambda: torch.nn.Linear(2, 2))
+
+    assert executor.model_for_training() is model
+    assert executor.model_for_inference() is model
+
+
+def test_fsdp_executor_rejects_online_rollout():
+    executor = FSDPExecutor()
+    assert not executor.rollout_capabilities().supported
+    with pytest.raises(RuntimeError, match="does not support online rollout"):
+        executor.model_for_inference()
+
+
 def test_spawn_only_on_rank():
     spawn_parallel_fn(only_on_rank, world_size=2, backend="gloo")
 
 
 def test_spawn_all_reduce():
     spawn_parallel_fn(all_reduce, world_size=2, backend="gloo")
+
+
+def test_spawn_ddp_model_views():
+    spawn_parallel_fn(
+        ddp_model_views,
+        world_size=2,
+        backend="gloo",
+        device_type="cpu",
+    )

tests/trainer/test_online_e2e.py

@@ -11,6 +11,7 @@
 from astrai.model.transformer import AutoRegressiveLM
 from astrai.trainer.rollout import BaseRewardModel
 from astrai.trainer.schedule import SchedulerFactory
+from astrai.trainer.train_context import TrainContextBuilder
 from astrai.trainer.trainer import Trainer
 from tests.helpers import CHAT_TEMPLATE
 
@@ -85,19 +86,26 @@ def _scheduler_fn(optim):
 ]
 
 
-@pytest.mark.integration
-@pytest.mark.parametrize(("strategy", "strategy_kwargs"), _ONLINE_STRATEGIES)
-def test_online_rollout_end_to_end(base_test_env, strategy, strategy_kwargs):
-    """Run one epoch of online RL rollout with KV-cache-backed generation."""
-    test_dir = base_test_env["test_dir"]
-    device = base_test_env["device"]
-    tokenizer = base_test_env["tokenizer"]
-    model_config = base_test_env["transformer_config"]
+_PARALLEL_CONFIGS = [
+    pytest.param("none", 1, "nccl", None, id="single"),
+    pytest.param("ddp", 1, "nccl", None, id="ddp-single"),
+    pytest.param("ddp", 2, "gloo", "cpu", id="ddp-cpu"),
+]
 
-    tokenizer.set_chat_template(CHAT_TEMPLATE)
-    tokenizer.save_pretrained(test_dir)
 
-    train_config = TrainConfig(
+def _make_train_config(
+    *,
+    test_dir,
+    model_config,
+    strategy,
+    strategy_kwargs,
+    parallel_mode="none",
+    nprocs=1,
+    backend="nccl",
+    device_type="cpu",
+    compile_mode=None,
+):
+    return TrainConfig(
         strategy=strategy,
         model_fn=partial(_model_fn, model_config),
         dataset=InstructionDataset(),
@@ -109,9 +117,11 @@ def test_online_rollout_end_to_end(base_test_env, strategy, strategy_kwargs):
         ckpt_interval=100,
         grad_accum_steps=1,
         random_seed=42,
-        device_type=device,
-        nprocs=1,
-        parallel_mode="none",
+        device_type=device_type,
+        nprocs=nprocs,
+        backend=backend,
+        parallel_mode=parallel_mode,
+        compile_mode=compile_mode,
         strategy_kwargs=strategy_kwargs,
         rollout_interval=1,
         rollout_temperature=1.0,
@@ -122,7 +132,75 @@ def test_online_rollout_end_to_end(base_test_env, strategy, strategy_kwargs):
         collate_fn=instruction_collate_fn,
     )
 
+
+@pytest.mark.integration
+@pytest.mark.parametrize(("strategy", "strategy_kwargs"), _ONLINE_STRATEGIES)
+@pytest.mark.parametrize(
+    ("parallel_mode", "nprocs", "backend", "runtime_device"), _PARALLEL_CONFIGS
+)
+def test_online_rollout_end_to_end(
+    base_test_env,
+    strategy,
+    strategy_kwargs,
+    parallel_mode,
+    nprocs,
+    backend,
+    runtime_device,
+):
+    """Run one epoch of online RL rollout with KV-cache-backed generation."""
+    test_dir = base_test_env["test_dir"]
+    device = base_test_env["device"]
+    tokenizer = base_test_env["tokenizer"]
+    model_config = base_test_env["transformer_config"]
+
+    tokenizer.set_chat_template(CHAT_TEMPLATE)
+    tokenizer.save_pretrained(test_dir)
+
+    train_config = _make_train_config(
+        test_dir=test_dir,
+        model_config=model_config,
+        strategy=strategy,
+        strategy_kwargs=strategy_kwargs,
+        parallel_mode=parallel_mode,
+        nprocs=nprocs,
+        backend=backend,
+        device_type=runtime_device or device,
+    )
+
     trainer = Trainer(train_config)
     trainer.train(param_path=test_dir)
 
     assert os.path.isdir(os.path.join(test_dir, "ckpt"))
+
+
+@pytest.mark.parametrize(
+    ("parallel_mode", "compile_mode", "match"),
+    [
+        pytest.param(
+            "fsdp",
+            None,
+            "Online rollout is not supported.*parallel_mode='fsdp'",
+            id="fsdp",
+        ),
+        pytest.param(
+            "none",
+            "default",
+            "Online rollout with torch.compile is not supported",
+            id="compile",
+        ),
+    ],
+)
+def test_online_rollout_fails_fast_for_unsupported_execution(
+    base_test_env, parallel_mode, compile_mode, match
+):
+    config = _make_train_config(
+        test_dir=base_test_env["test_dir"],
+        model_config=base_test_env["transformer_config"],
+        strategy="online_grpo",
+        strategy_kwargs={"group_size": 2},
+        parallel_mode=parallel_mode,
+        compile_mode=compile_mode,
+    )
+
+    with pytest.raises(ValueError, match=match):
+        TrainContextBuilder(config).build()