diff --git a/astrai/parallel/__init__.py b/astrai/parallel/__init__.py index b565f8d4..fffdfc2d 100644 --- a/astrai/parallel/__init__.py +++ b/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", ] diff --git a/astrai/parallel/executor.py b/astrai/parallel/executor.py index d0c12c86..a77984d4 100644 --- a/astrai/parallel/executor.py +++ b/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") diff --git a/astrai/trainer/train_context.py b/astrai/trainer/train_context.py index e41a695e..fac31885 100644 --- a/astrai/trainer/train_context.py +++ b/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, diff --git a/tests/parallel/test_parallel.py b/tests/parallel/test_parallel.py index c05b0848..a27251f0 100644 --- a/tests/parallel/test_parallel.py +++ b/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", + ) diff --git a/tests/trainer/test_online_e2e.py b/tests/trainer/test_online_e2e.py index 9f1eac86..3c2d2e32 100644 --- a/tests/trainer/test_online_e2e.py +++ b/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()