fix: keep async rollouts version-consistent

- serialize shared-model optimizer updates with generation
- reject future or over-lagged rollout results after asynchronous scoring
- close cache publication races
- persist policy versions in online checkpoints

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:

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 _validate_weight_version(
-    def update_weights(self, policy_version: int) -> int:
+        self, policy_version: int, *, require_advance: bool = False
-        """Acknowledge an in-place weight update and invalidate stale KV state.
+    ) -> None:
-
-        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.
-        """
         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:
+        if require_advance and policy_version == self._policy_version:
-            return 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)
 

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._validate_policy_version(raw)
-            self._cache_key = cache_key
+            scored = self._score(raw)
-            self._steps_since_rollout = 0
+
-            return self._cache, True
+            def commit(live_version: int) -> Tuple[RolloutResult, bool]:
-        return self._cache, False
+                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)

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)
+            raise RuntimeError(
-            self._rollout_runner.step()
+                "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."""

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,

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,
             )
         )

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.optimizer_step(context.optimizer)
-                            context.strategy.on_optimizer_step()
                             context.optimizer.zero_grad()
 
                             if context.scheduler:

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]

docs/developer/internals.md

@@ -119,8 +119,7 @@ on_train_begin
 
         if executor.sync_gradients:
           before_optimizer_step
-          optimizer.step()
+          strategy.optimizer_step(optimizer)
-          strategy.on_optimizer_step()
           optimizer.zero_grad()
           if scheduler:
             scheduler.step()

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 |

docs/guides/training.md

@@ -71,8 +71,7 @@ on_train_begin
 
         if executor.sync_gradients:
           before_optimizer_step
-          optimizer.step()
+          strategy.optimizer_step(optimizer)
-          strategy.on_optimizer_step()
           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
+Every successful optimizer step mutates the shared model and advances its
-acknowledges the shared-model weight update to the rollout scheduler. The
+monotonic `policy_version` under the same generation lock. The scheduler
-scheduler invalidates reusable KV prefixes before accepting the new version.
+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

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,

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")

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

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)))

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))