diff --git a/astrai/config/train_config.py b/astrai/config/train_config.py index 4b2dc54..6a62be9 100644 --- a/astrai/config/train_config.py +++ b/astrai/config/train_config.py @@ -64,6 +64,7 @@ class TrainConfig(BaseConfig): neftune_alpha (float): NEFTune noise alpha, 0=disabled, typical: 5.0. Defaults to 0.0. moe_aux_loss_coef (float): Weight applied to the MoE load-balancing loss. Defaults to 0.01. rollout_interval (int): Number of optimizer steps between online rollouts. Defaults to 512. + rollout_max_policy_lag (Optional[int]): Maximum accepted gap between rollout and live policy versions. None derives ``rollout_interval - 1``. Defaults to None. rollout_temperature (float): Sampling temperature for online rollout. Defaults to 0.7. rollout_top_k (int): Top-k filtering for online rollout, 0=disable. Defaults to 0. rollout_top_p (float): Top-p (nucleus) filtering for online rollout. Defaults to 0.9. @@ -118,6 +119,7 @@ class TrainConfig(BaseConfig): moe_aux_loss_coef: float = 0.01 rollout_interval: int = 512 + rollout_max_policy_lag: Optional[int] = None rollout_temperature: float = 0.7 rollout_top_k: int = 0 rollout_top_p: float = 0.9 @@ -199,6 +201,12 @@ class TrainConfig(BaseConfig): raise ValueError(f"must be non-negative, got {v}") return v + @field_validator("rollout_max_policy_lag") + def _validate_optional_non_negative_int(cls, v: Optional[int]) -> Optional[int]: + if v is not None and v < 0: + raise ValueError(f"rollout_max_policy_lag must be non-negative, got {v}") + return v + @field_validator("max_grad_norm") def _validate_max_grad_norm(cls, v: Optional[float]) -> Optional[float]: if v is not None and v <= 0: diff --git a/astrai/inference/scheduler.py b/astrai/inference/scheduler.py index 3121ae3..3de9ac0 100644 --- a/astrai/inference/scheduler.py +++ b/astrai/inference/scheduler.py @@ -3,7 +3,7 @@ import threading import uuid from contextlib import nullcontext from functools import wraps -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any, Callable, Dict, List, Optional, Tuple, TypeVar, Union import torch @@ -27,6 +27,7 @@ from astrai.model.automodel import AutoModel from astrai.tokenize.tokenizer import AutoTokenizer logger = logging.getLogger(__name__) +T = TypeVar("T") def _with_weight_lock(method): @@ -128,15 +129,9 @@ class InferenceScheduler: """Version of the model weights used for subsequent generations.""" return self._policy_version - @_with_weight_lock - def update_weights(self, policy_version: int) -> int: - """Acknowledge an in-place weight update and invalidate stale KV state. - - The scheduler owns the same model object as the in-process trainer, so - weights have already changed when this method is called. The explicit - version update makes that lifecycle visible and prevents prefix KV - entries produced by older weights from being reused. - """ + def _validate_weight_version( + self, policy_version: int, *, require_advance: bool = False + ) -> None: if ( isinstance(policy_version, bool) or not isinstance(policy_version, int) @@ -148,17 +143,57 @@ class InferenceScheduler: f"policy_version cannot move backwards from " f"{self._policy_version} to {policy_version}" ) - if policy_version == self._policy_version: - return self._policy_version + if require_advance and policy_version == self._policy_version: + raise ValueError( + f"policy_version must advance beyond {self._policy_version} " + "when model weights are mutated" + ) + + def _ensure_weight_update_ready(self) -> None: if self._loop_thread is not None and self._loop_thread.is_alive(): raise RuntimeError("Stop the scheduler before updating model weights") if self._task_mgr.get_active_tasks() or self._task_mgr.get_waiting_tasks(): raise RuntimeError("Cannot update model weights while tasks are queued") + def _commit_weight_version(self, policy_version: int) -> int: self._task_cache.invalidate_cache() self._policy_version = policy_version return self._policy_version + @_with_weight_lock + def update_weights(self, policy_version: int) -> int: + """Acknowledge an in-place weight update and invalidate stale KV state. + + The scheduler owns the same model object as the in-process trainer, so + weights have already changed when this method is called. The explicit + version update makes that lifecycle visible and prevents prefix KV + entries produced by older weights from being reused. + """ + self._validate_weight_version(policy_version) + if policy_version == self._policy_version: + return self._policy_version + self._ensure_weight_update_ready() + return self._commit_weight_version(policy_version) + + @_with_weight_lock + def apply_weight_update(self, policy_version: int, update: Callable[[], T]) -> T: + """Mutate shared weights and publish their version without generation.""" + if not callable(update): + raise TypeError("update must be callable") + self._validate_weight_version(policy_version, require_advance=True) + self._ensure_weight_update_ready() + + result = update() + self._commit_weight_version(policy_version) + return result + + @_with_weight_lock + def with_policy_snapshot(self, inspect: Callable[[int], T]) -> T: + """Inspect state while the scheduler's policy version remains stable.""" + if not callable(inspect): + raise TypeError("inspect must be callable") + return inspect(self._policy_version) + def add_task(self, prompt: str, **kwargs) -> str: return self._task_mgr.add_task(prompt, **kwargs) diff --git a/astrai/trainer/rollout.py b/astrai/trainer/rollout.py index efa0456..7e81d2c 100644 --- a/astrai/trainer/rollout.py +++ b/astrai/trainer/rollout.py @@ -16,7 +16,7 @@ Provides: import threading from abc import ABC, abstractmethod from dataclasses import dataclass, field -from typing import Dict, List, Optional, Tuple +from typing import Callable, Dict, List, Optional, Tuple, TypeVar import torch from torch import Tensor @@ -98,6 +98,11 @@ class BaseRewardModel(ABC): _PAD = 0 +T = TypeVar("T") + + +class RolloutVersionError(RuntimeError): + """A rollout cannot be attributed to an acceptable policy version.""" class RolloutGenerator: @@ -142,6 +147,18 @@ class RolloutGenerator: with self._weight_lock: return self.scheduler.update_weights(policy_version) + def apply_weight_update(self, policy_version: int, update: Callable[[], T]) -> T: + """Apply a shared-model mutation at an atomic generation boundary.""" + with self._weight_lock: + return self.scheduler.apply_weight_update(policy_version, update) + + def with_policy_snapshot(self, inspect: Callable[[int], T]) -> T: + """Inspect a version stable against generator and scheduler updates.""" + if not callable(inspect): + raise TypeError("inspect must be callable") + with self._weight_lock: + return self.scheduler.with_policy_snapshot(inspect) + @torch.no_grad() def generate(self, batch: Dict) -> RawRollout: """Expand prompts by ``group_size`` and generate one response each. @@ -159,15 +176,22 @@ class RolloutGenerator: format the policy was SFT-trained on. """ with self._weight_lock: - model = self.scheduler._executor.model - was_training = model.training - model.eval() - try: - return self._generate_eval(batch) - finally: - model.train(was_training) - def _generate_eval(self, batch: Dict) -> RawRollout: + def generate_snapshot(generation_version: int) -> RawRollout: + model = self.scheduler._executor.model + was_training = model.training + model.eval() + try: + return self._generate_eval(batch, generation_version) + finally: + model.train(was_training) + + # Capture the version under the scheduler lock as well as the + # generator lock. This also serializes callers that update the + # scheduler directly instead of going through this wrapper. + return self.scheduler.with_policy_snapshot(generate_snapshot) + + def _generate_eval(self, batch: Dict, generation_version: int) -> RawRollout: prompt_texts, flat_prompt_ids = self._prepare_prompts(batch) B = len(prompt_texts) G = self.group_size @@ -258,7 +282,7 @@ class RolloutGenerator: responses=responses, response_mask=response_mask, logprobs_old=logprobs_old, - policy_version=self.policy_version, + policy_version=generation_version, prompt_texts=prompt_texts, response_texts=response_texts, ) @@ -375,10 +399,18 @@ class RolloutRunner: generator: RolloutGenerator, reward_model: BaseRewardModel, rollout_interval: int = 512, + max_policy_lag: Optional[int] = None, ): + if rollout_interval <= 0: + raise ValueError("rollout_interval must be positive") + if max_policy_lag is not None and max_policy_lag < 0: + raise ValueError("max_policy_lag must be non-negative or None") self.generator = generator self.reward_model = reward_model self.rollout_interval = rollout_interval + self.max_policy_lag = ( + rollout_interval - 1 if max_policy_lag is None else max_policy_lag + ) self._cache: Optional[RolloutResult] = None self._cache_key = None @@ -392,6 +424,10 @@ class RolloutRunner: """Publish the shared policy's new version to the rollout backend.""" return self.generator.update_weights(policy_version) + def apply_weight_update(self, policy_version: int, update: Callable[[], T]) -> T: + """Apply a model update and publish its version as one operation.""" + return self.generator.apply_weight_update(policy_version, update) + def step(self): """Advance the internal counter (call once per optimizer step).""" self._steps_since_rollout += 1 @@ -442,6 +478,26 @@ class RolloutRunner: response_texts=raw.response_texts, ) + def _validate_policy_version( + self, result: RawRollout, *, live_version: Optional[int] = None + ) -> None: + version = result.policy_version + if isinstance(version, bool) or not isinstance(version, int) or version < 0: + raise RolloutVersionError(f"rollout has invalid policy version {version!r}") + if live_version is None: + live_version = self.policy_version + if version > live_version: + raise RolloutVersionError( + f"rollout has future policy version {version}; " + f"live policy version is {live_version}" + ) + lag = live_version - version + if lag > self.max_policy_lag: + raise RolloutVersionError( + f"rollout policy lag {lag} exceeds max_policy_lag=" + f"{self.max_policy_lag} (rollout={version}, live={live_version})" + ) + def __call__(self, batch: Dict[str, Tensor]) -> Tuple[RolloutResult, bool]: """Return ``(cached or fresh) RolloutResult`` plus an ``is_fresh`` flag. @@ -455,8 +511,26 @@ class RolloutRunner: or self._steps_since_rollout >= self.rollout_interval ): raw = self.generator.generate(batch) - self._cache = self._score(raw) - self._cache_key = cache_key - self._steps_since_rollout = 0 - return self._cache, True - return self._cache, False + self._validate_policy_version(raw) + scored = self._score(raw) + + def commit(live_version: int) -> Tuple[RolloutResult, bool]: + self._validate_policy_version(scored, live_version=live_version) + self._cache = scored + self._cache_key = cache_key + self._steps_since_rollout = 0 + return scored, True + + # A weight update cannot land between the final version check and + # cache publication. Reward scoring itself intentionally remains + # outside the policy lock because it may call an external service. + return self.generator.with_policy_snapshot(commit) + + cached = self._cache + assert cached is not None + + def reuse(live_version: int) -> Tuple[RolloutResult, bool]: + self._validate_policy_version(cached, live_version=live_version) + return cached, False + + return self.generator.with_policy_snapshot(reuse) diff --git a/astrai/trainer/strategy.py b/astrai/trainer/strategy.py index 2321c7d..538fc79 100644 --- a/astrai/trainer/strategy.py +++ b/astrai/trainer/strategy.py @@ -7,6 +7,7 @@ import torch import torch.nn as nn import torch.nn.functional as F from torch import Tensor +from torch.optim import Optimizer from astrai.factory import BaseFactory from astrai.model.components.mlp import RouterStats @@ -279,10 +280,22 @@ class BaseStrategy(ABC): self._moe_metrics["aux_loss"] = float(aux_loss.detach().cpu().item()) def on_optimizer_step(self): - """Advance online rollout state after a successful optimizer step.""" + """Reject unsafe post-hoc publication for an online shared model.""" if self._rollout_runner is not None: - self._rollout_runner.update_weights(self.policy_version + 1) - self._rollout_runner.step() + raise RuntimeError( + "online training must call strategy.optimizer_step(optimizer) " + "so weight mutation and policy-version publication are atomic" + ) + + def optimizer_step(self, optimizer: Optimizer): + """Step the optimizer at an atomic online-rollout version boundary.""" + if self._rollout_runner is None: + return optimizer.step() + + next_version = self.policy_version + 1 + result = self._rollout_runner.apply_weight_update(next_version, optimizer.step) + self._rollout_runner.step() + return result def __call__(self, batch: Dict[str, Tensor]) -> LossOutput: """Run offline or online forward depending on runner injection.""" diff --git a/astrai/trainer/train_callback.py b/astrai/trainer/train_callback.py index ba5c1ef..a1c8677 100644 --- a/astrai/trainer/train_callback.py +++ b/astrai/trainer/train_callback.py @@ -164,6 +164,9 @@ class CheckpointCallback(TrainCallback): **context.config.to_dict(), "optimizer_step": context.optimizer_step, } + policy_version = context.strategy.policy_version + if policy_version is not None: + meta["policy_version"] = policy_version context.checkpoint = Checkpoint( state_dict=state_dict, epoch=context.epoch, diff --git a/astrai/trainer/train_context.py b/astrai/trainer/train_context.py index ff6ac9a..5bac667 100644 --- a/astrai/trainer/train_context.py +++ b/astrai/trainer/train_context.py @@ -355,5 +355,6 @@ class TrainContextBuilder: generator=generator, reward_model=cfg.reward_model_fn(), rollout_interval=cfg.rollout_interval, + max_policy_lag=cfg.rollout_max_policy_lag, ) ) diff --git a/astrai/trainer/trainer.py b/astrai/trainer/trainer.py index 09fcb3a..6937574 100644 --- a/astrai/trainer/trainer.py +++ b/astrai/trainer/trainer.py @@ -94,8 +94,7 @@ class Trainer: if executor.sync_gradients: self._call_callbacks("before_optimizer_step", context) - context.optimizer.step() - context.strategy.on_optimizer_step() + context.strategy.optimizer_step(context.optimizer) context.optimizer.zero_grad() if context.scheduler: diff --git a/docs/developer/architecture.md b/docs/developer/architecture.md index 50358a8..7a4fba7 100644 --- a/docs/developer/architecture.md +++ b/docs/developer/architecture.md @@ -625,7 +625,7 @@ classDiagram +supports_online() bool +set_rollout_runner(runner) +prepare_from_rollout(result) Dict - +on_optimizer_step() + +optimizer_step(optimizer) } class LossOutput { @@ -698,12 +698,15 @@ classDiagram +int rep_window +int policy_version +update_weights(policy_version) int + +apply_weight_update(policy_version, update) +generate(batch) RawRollout } class RolloutRunner { +int policy_version + +int max_policy_lag +update_weights(policy_version) int + +apply_weight_update(policy_version, update) +step() +clear_cache() +__call__(batch) Tuple[RolloutResult, bool] diff --git a/docs/developer/internals.md b/docs/developer/internals.md index ccfdd8f..b884964 100644 --- a/docs/developer/internals.md +++ b/docs/developer/internals.md @@ -119,8 +119,7 @@ on_train_begin if executor.sync_gradients: before_optimizer_step - optimizer.step() - strategy.on_optimizer_step() + strategy.optimizer_step(optimizer) optimizer.zero_grad() if scheduler: scheduler.step() diff --git a/docs/guides/params.md b/docs/guides/params.md index 9b23ebd..7dd8ce8 100644 --- a/docs/guides/params.md +++ b/docs/guides/params.md @@ -156,6 +156,7 @@ provide a command-line option for configuring one. | Parameter | Description | Default | |-----------|-------------|---------| | `--rollout_interval` | Optimizer steps between rollout refreshes | 512 | +| `--rollout_max_policy_lag` | Maximum accepted rollout/live policy-version gap (`None` derives `rollout_interval - 1`) | None | | `--rollout_temperature` | Rollout sampling temperature | 0.7 | | `--rollout_top_k` | Rollout top-k filtering (`0` disables) | 0 | | `--rollout_top_p` | Rollout nucleus sampling threshold | 0.9 | diff --git a/docs/guides/training.md b/docs/guides/training.md index 0d267f1..bba80aa 100644 --- a/docs/guides/training.md +++ b/docs/guides/training.md @@ -71,8 +71,7 @@ on_train_begin if executor.sync_gradients: before_optimizer_step - optimizer.step() - strategy.on_optimizer_step() + strategy.optimizer_step(optimizer) optimizer.zero_grad() if scheduler: scheduler.step() @@ -171,12 +170,16 @@ them with a `BaseRewardModel`. It refreshes cached rollouts every behaviour log-probabilities into the loss, so it does not allocate or synchronize a separate old-policy model. -Every successful optimizer step advances a monotonic `policy_version` and -acknowledges the shared-model weight update to the rollout scheduler. The -scheduler invalidates reusable KV prefixes before accepting the new version. +Every successful optimizer step mutates the shared model and advances its +monotonic `policy_version` under the same generation lock. The scheduler +invalidates reusable KV prefixes before accepting the new version, so an async +rollout cannot observe partially updated weights under the previous version. `RawRollout` and `RolloutResult` retain the version that actually generated their behavior log-probabilities, so cached rollout samples remain attributable -even while later optimizer steps advance the live policy. +even while later optimizer steps advance the live policy. Results from a future +version or beyond `rollout_max_policy_lag` are rejected before training. The +final version check and rollout-cache publication share that policy lock, so a +concurrent update cannot land between validation and cache insertion. Online strategies require `TrainConfig.reward_model_fn`. `train.py` exposes the rollout sampling parameters but does not yet offer a CLI argument for the reward diff --git a/scripts/tools/train.py b/scripts/tools/train.py index 1f9690e..2f15631 100644 --- a/scripts/tools/train.py +++ b/scripts/tools/train.py @@ -333,6 +333,13 @@ _START_METHODS = sorted(START_METHODS) group="Algorithm", help="Steps between rollouts.", ) +@opt( + "--rollout_max_policy_lag", + type=int, + default=None, + group="Algorithm", + help="Maximum accepted rollout/live policy-version gap.", +) @opt( "--rollout_temperature", type=float, @@ -684,6 +691,7 @@ def train( } rollout_interval = kwargs.pop("rollout_interval", 512) + rollout_max_policy_lag = kwargs.pop("rollout_max_policy_lag", None) rollout_temperature = kwargs.pop("rollout_temperature", 0.7) rollout_top_k = kwargs.pop("rollout_top_k", 0) rollout_top_p = kwargs.pop("rollout_top_p", 0.9) @@ -840,6 +848,7 @@ def train( neftune_alpha=neftune_alpha, collate_fn=collate_fn, rollout_interval=rollout_interval, + rollout_max_policy_lag=rollout_max_policy_lag, rollout_temperature=rollout_temperature, rollout_top_k=rollout_top_k, rollout_top_p=rollout_top_p, diff --git a/tests/inference/test_scheduler.py b/tests/inference/test_scheduler.py index 8879813..292b3cd 100644 --- a/tests/inference/test_scheduler.py +++ b/tests/inference/test_scheduler.py @@ -561,6 +561,78 @@ def test_scheduler_weight_versions_are_monotonic_and_acknowledged(device): scheduler.stop() +def test_scheduler_applies_weight_mutation_and_version_atomically(device): + scheduler, _tok, model = _make_real_scheduler(device) + before = next(model.parameters()).detach().clone() + + def mutate(): + with torch.no_grad(): + next(model.parameters()).add_(1) + return "updated" + + try: + assert scheduler.apply_weight_update(1, mutate) == "updated" + assert scheduler.policy_version == 1 + assert not torch.equal(next(model.parameters()), before) + with pytest.raises(ValueError, match="must advance"): + scheduler.apply_weight_update(1, mutate) + + def failed_mutation(): + raise RuntimeError("optimizer failed") + + with pytest.raises(RuntimeError, match="optimizer failed"): + scheduler.apply_weight_update(2, failed_mutation) + assert scheduler.policy_version == 1 + finally: + scheduler.stop() + + +def test_scheduler_serializes_policy_snapshot_and_direct_update(device): + scheduler, _tok, _model = _make_real_scheduler(device) + snapshot_started = threading.Event() + release_snapshot = threading.Event() + update_finished = threading.Event() + errors = [] + + def inspect(version): + assert version == 0 + snapshot_started.set() + assert release_snapshot.wait(timeout=5) + + def take_snapshot(): + try: + scheduler.with_policy_snapshot(inspect) + except BaseException as exc: + errors.append(exc) + + def update(): + try: + scheduler.update_weights(1) + update_finished.set() + except BaseException as exc: + errors.append(exc) + + snapshot_thread = threading.Thread(target=take_snapshot) + update_thread = threading.Thread(target=update) + try: + snapshot_thread.start() + assert snapshot_started.wait(timeout=5) + update_thread.start() + assert not update_finished.wait(timeout=0.1) + release_snapshot.set() + snapshot_thread.join(timeout=5) + update_thread.join(timeout=5) + assert not snapshot_thread.is_alive() + assert not update_thread.is_alive() + assert errors == [] + assert scheduler.policy_version == 1 + finally: + release_snapshot.set() + snapshot_thread.join(timeout=5) + update_thread.join(timeout=5) + scheduler.stop() + + def test_scheduler_rejects_weight_update_with_queued_tasks(device): scheduler, _tok, _model = _make_real_scheduler(device) task_id = scheduler.add_task("queued") diff --git a/tests/trainer/test_online_e2e.py b/tests/trainer/test_online_e2e.py index df8b817..9e0aaca 100644 --- a/tests/trainer/test_online_e2e.py +++ b/tests/trainer/test_online_e2e.py @@ -10,6 +10,7 @@ from torch.utils.data import Dataset import astrai.trainer.train_context as train_context from astrai.config import TrainConfig from astrai.model.transformer import AutoRegressiveLM +from astrai.serialization import Checkpoint from astrai.trainer.rollout import BaseRewardModel from astrai.trainer.schedule import SchedulerFactory from astrai.trainer.trainer import Trainer @@ -126,6 +127,7 @@ def test_online_rollout_end_to_end( parallel_mode="none", strategy_kwargs=strategy_kwargs, rollout_interval=1, + rollout_max_policy_lag=0, rollout_temperature=1.0, rollout_top_k=0, rollout_top_p=1.0, @@ -137,5 +139,8 @@ def test_online_rollout_end_to_end( trainer = Trainer(train_config) trainer.train(param_path=test_dir) - assert os.path.isdir(os.path.join(test_dir, "ckpt")) + checkpoint_dir = os.path.join(test_dir, "ckpt", "epoch_0_step_2") + assert os.path.isdir(checkpoint_dir) + checkpoint = Checkpoint.load(checkpoint_dir) + assert checkpoint.meta["policy_version"] == 2 assert len(created_reference_models) == 1 diff --git a/tests/trainer/test_online_strategy.py b/tests/trainer/test_online_strategy.py index e66a500..396231f 100644 --- a/tests/trainer/test_online_strategy.py +++ b/tests/trainer/test_online_strategy.py @@ -60,11 +60,25 @@ class _RecordingRunner: self.weight_updates.append(policy_version) return policy_version + def apply_weight_update(self, policy_version, update): + result = update() + self.update_weights(policy_version) + return result + def swap_result(self, result): self.result = result self._fresh = True +class _NoOpOptimizer: + def step(self): + return None + + +def _step(strat): + strat.optimizer_step(_NoOpOptimizer()) + + def _make_grpo(device, executor=None): model, _ = make_model(device) ref_model = make_frozen(model, device) @@ -250,9 +264,9 @@ def test_grpo_reuses_same_cached_result(device): runner = _RecordingRunner(_make_rollout_result(device=device)) strat.set_rollout_runner(runner) strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)}) - strat.on_optimizer_step() + _step(strat) strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)}) - strat.on_optimizer_step() + _step(strat) assert runner.calls == 2 assert runner.step_calls == 2 @@ -262,10 +276,10 @@ def test_grpo_accepts_new_rollout_result(device): runner = _RecordingRunner(_make_rollout_result(device=device)) strat.set_rollout_runner(runner) strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)}) - strat.on_optimizer_step() + _step(strat) runner.swap_result(_make_rollout_result(device=device)) strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)}) - strat.on_optimizer_step() + _step(strat) assert runner.calls == 2 assert runner.step_calls == 2 @@ -280,10 +294,10 @@ def test_dpo_no_sync_hook_when_new_rollout_result(device): runner = _RecordingRunner(_make_rollout_result(device=device)) strat.set_rollout_runner(runner) strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)}) - strat.on_optimizer_step() + _step(strat) runner.swap_result(_make_rollout_result(device=device)) strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)}) - strat.on_optimizer_step() + _step(strat) assert runner.step_calls == 2 @@ -302,12 +316,36 @@ def test_step_called_when_sync_gradients_true(device): runner = _RecordingRunner(_make_rollout_result(device=device)) strat.set_rollout_runner(runner) strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)}) - strat.on_optimizer_step() + _step(strat) assert runner.step_calls == 1 assert runner.weight_updates == [1] assert strat.policy_version == 1 +def test_post_hoc_online_optimizer_step_is_rejected(device): + strat = _make_grpo(device) + strat.set_rollout_runner(_RecordingRunner(_make_rollout_result(device=device))) + + with pytest.raises(RuntimeError, match="strategy.optimizer_step"): + strat.on_optimizer_step() + + +def test_optimizer_step_publishes_version_with_weight_update(device): + strat = _make_grpo(device) + runner = _RecordingRunner(_make_rollout_result(device=device)) + strat.set_rollout_runner(runner) + parameter = next(strat.model.parameters()) + parameter.grad = torch.ones_like(parameter) + optimizer = torch.optim.SGD(strat.model.parameters(), lr=0.1) + before = parameter.detach().clone() + + strat.optimizer_step(optimizer) + + assert not torch.equal(parameter, before) + assert runner.weight_updates == [1] + assert runner.step_calls == 1 + + def test_loss_is_differentiable_dpo(device): strat = _make_dpo(device) strat.set_rollout_runner(_RecordingRunner(_make_rollout_result(device=device))) diff --git a/tests/trainer/test_rollout.py b/tests/trainer/test_rollout.py index a36ad6d..e22c04a 100644 --- a/tests/trainer/test_rollout.py +++ b/tests/trainer/test_rollout.py @@ -1,5 +1,7 @@ """Unit tests for the online rollout module.""" +import threading + import pytest import torch @@ -11,6 +13,7 @@ from astrai.trainer.rollout import ( RolloutGenerator, RolloutResult, RolloutRunner, + RolloutVersionError, ) from tests.helpers import FakeTokenizer, make_model @@ -151,6 +154,113 @@ def test_rollout_generator_uses_eval_and_restores_mode(device): assert model.training is True +def test_rollout_generator_serializes_generation_and_policy_update(device): + gen, _ = _make_generator(device, group_size=1, max_tokens=2) + generation_started = threading.Event() + allow_generation_to_finish = threading.Event() + update_finished = threading.Event() + thread_errors = [] + original = gen._generate_eval + + def blocking_generate(batch, generation_version): + generation_started.set() + assert allow_generation_to_finish.wait(timeout=5) + return original(batch, generation_version) + + gen._generate_eval = blocking_generate + + def generate(): + try: + gen.generate(_make_instruction_batch(n=1)) + except BaseException as exc: + thread_errors.append(exc) + + def apply_update(): + try: + gen.apply_weight_update(1, update_finished.set) + except BaseException as exc: + thread_errors.append(exc) + + generation_thread = threading.Thread(target=generate) + update_thread = threading.Thread(target=apply_update) + generation_thread.start() + assert generation_started.wait(timeout=5) + update_thread.start() + assert not update_finished.wait(timeout=0.1) + + allow_generation_to_finish.set() + generation_thread.join(timeout=5) + update_thread.join(timeout=5) + assert not generation_thread.is_alive() + assert not update_thread.is_alive() + assert thread_errors == [] + assert update_finished.is_set() + assert gen.policy_version == 1 + + +def test_rollout_generator_serializes_direct_scheduler_update(device): + gen, _ = _make_generator(device, group_size=1, max_tokens=2) + generation_started = threading.Event() + allow_generation_to_finish = threading.Event() + update_finished = threading.Event() + thread_errors = [] + original = gen._generate_eval + + def blocking_generate(batch, generation_version): + generation_started.set() + assert allow_generation_to_finish.wait(timeout=5) + return original(batch, generation_version) + + gen._generate_eval = blocking_generate + rollout = [] + + def generate(): + try: + rollout.append(gen.generate(_make_instruction_batch(n=1))) + except BaseException as exc: + thread_errors.append(exc) + + def update_scheduler_directly(): + try: + gen.scheduler.update_weights(1) + update_finished.set() + except BaseException as exc: + thread_errors.append(exc) + + generation_thread = threading.Thread(target=generate) + update_thread = threading.Thread(target=update_scheduler_directly) + generation_thread.start() + assert generation_started.wait(timeout=5) + update_thread.start() + assert not update_finished.wait(timeout=0.1) + + allow_generation_to_finish.set() + generation_thread.join(timeout=5) + update_thread.join(timeout=5) + assert not generation_thread.is_alive() + assert not update_thread.is_alive() + assert thread_errors == [] + assert rollout[0].policy_version == 0 + assert gen.policy_version == 1 + + +def test_rollout_generator_keeps_generation_start_version(device): + gen, _ = _make_generator(device, group_size=1, max_tokens=2) + original_run_batch = gen.scheduler.run_batch + + def update_after_generation(*args, **kwargs): + result = original_run_batch(*args, **kwargs) + gen.scheduler.update_weights(1) + return result + + gen.scheduler.run_batch = update_after_generation + + rollout = gen.generate(_make_instruction_batch(n=1)) + + assert rollout.policy_version == 0 + assert gen.policy_version == 1 + + def test_rollout_generator_mask_matches_responses(device): """Positions beyond a response's length are pad (mask False).""" gen, _ = _make_generator(device, group_size=2, max_tokens=6) @@ -248,6 +358,7 @@ def _make_runner(device, **kw): generator=generator, reward_model=rm, rollout_interval=kw.get("rollout_interval", 2), + max_policy_lag=kw.get("max_policy_lag"), ), model, ) @@ -297,6 +408,114 @@ def test_rollout_runner_tags_generation_version_and_preserves_cached_behavior(de assert refreshed.policy_version == 1 +def test_rollout_runner_rejects_future_generation_version(device): + runner, _ = _make_runner(device, rollout_interval=2) + raw = runner.generator.generate(_make_instruction_batch(n=1)) + raw.policy_version = runner.policy_version + 1 + runner.generator.generate = lambda _batch: raw + + with pytest.raises(RolloutVersionError, match="future policy version"): + runner(_make_instruction_batch(n=1)) + + +def test_rollout_runner_rejects_result_beyond_max_policy_lag(device): + runner, _ = _make_runner(device, rollout_interval=4, max_policy_lag=1) + batch = _make_instruction_batch(n=1) + result, _ = runner(batch) + assert result.policy_version == 0 + + runner.update_weights(2) + with pytest.raises(RolloutVersionError, match="exceeds max_policy_lag=1"): + runner(batch) + + +def test_rollout_runner_revalidates_version_after_async_scoring(device): + runner, _ = _make_runner(device, rollout_interval=4, max_policy_lag=0) + original_score = runner._score + + def score_while_policy_advances(raw): + result = original_score(raw) + runner.update_weights(1) + return result + + runner._score = score_while_policy_advances + + with pytest.raises(RolloutVersionError, match="exceeds max_policy_lag=0"): + runner(_make_instruction_batch(n=1)) + assert runner._cache is None + + +def test_rollout_runner_publishes_cache_before_concurrent_policy_update(device): + runner, _ = _make_runner(device, rollout_interval=4, max_policy_lag=1) + final_validation_started = threading.Event() + allow_final_validation_to_finish = threading.Event() + update_finished = threading.Event() + rollout_finished = threading.Event() + thread_errors = [] + validation_calls = 0 + original_validate = runner._validate_policy_version + + def blocking_validate(result, *, live_version=None): + nonlocal validation_calls + validation_calls += 1 + original_validate(result, live_version=live_version) + if validation_calls == 2: + final_validation_started.set() + assert allow_final_validation_to_finish.wait(timeout=5) + + runner._validate_policy_version = blocking_validate + + def produce_rollout(): + try: + runner(_make_instruction_batch(n=1)) + rollout_finished.set() + except BaseException as exc: + thread_errors.append(exc) + + def apply_update(): + try: + runner.apply_weight_update(1, update_finished.set) + except BaseException as exc: + thread_errors.append(exc) + + rollout_thread = threading.Thread(target=produce_rollout) + update_thread = threading.Thread(target=apply_update) + rollout_thread.start() + assert final_validation_started.wait(timeout=5) + update_thread.start() + assert not update_finished.wait(timeout=0.1) + + allow_final_validation_to_finish.set() + rollout_thread.join(timeout=5) + update_thread.join(timeout=5) + assert not rollout_thread.is_alive() + assert not update_thread.is_alive() + assert thread_errors == [] + assert rollout_finished.is_set() + assert update_finished.is_set() + assert runner._cache is not None + assert runner._cache.policy_version == 0 + assert runner.policy_version == 1 + + +def test_rollout_runner_derives_default_policy_lag_from_interval(device): + runner, _ = _make_runner(device, rollout_interval=4) + assert runner.max_policy_lag == 3 + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"rollout_interval": 0}, "rollout_interval must be positive"), + ({"max_policy_lag": -1}, "max_policy_lag must be non-negative"), + ], +) +def test_rollout_runner_rejects_invalid_version_window(device, kwargs, message): + generator, _ = _make_generator(device) + with pytest.raises(ValueError, match=message): + RolloutRunner(generator, ConstantRewardModel(), **kwargs) + + def test_rollout_runner_refreshes_for_different_batch(device): runner, _ = _make_runner(device, rollout_interval=100) r1, fresh1 = runner(_make_instruction_batch(n=1))