fix: keep async rollouts version-consistent

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 @@ def _validate_non_negative(cls, v):
             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 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.tokenize.tokenizer import AutoTokenizer
 
 logger = logging.getLogger(__name__)
+T = TypeVar("T")
 
 
 def _with_weight_lock(method):
@@ -128,15 +129,9 @@ def policy_version(self) -> int:
         """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 @@ def update_weights(self, policy_version: int) -> int:
                 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)
 

astrai/trainer/rollout.py

@@ -16,7 +16,7 @@
 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 @@ def score(self, prompts: List[str], responses: List[List[str]]) -> Tensor:
 
 
 _PAD = 0
+T = TypeVar("T")
+
+
+class RolloutVersionError(RuntimeError):
+    """A rollout cannot be attributed to an acceptable policy version."""
 
 
 class RolloutGenerator:
@@ -142,6 +147,18 @@ def update_weights(self, policy_version: int) -> int:
         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 @@ def generate(self, batch: Dict) -> RawRollout:
         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 @@ def _generate_eval(self, batch: Dict) -> RawRollout:
             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 @@ def __init__(
         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 @@ def update_weights(self, policy_version: int) -> int:
         """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 @@ def _score(self, raw: RawRollout) -> RolloutResult:
             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 @@ def __call__(self, batch: Dict[str, Tensor]) -> Tuple[RolloutResult, bool]:
             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)

astrai/trainer/strategy.py

@@ -7,6 +7,7 @@
 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 @@ def _refresh_moe_diagnostics(
         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."""

astrai/trainer/train_callback.py

@@ -164,6 +164,9 @@ def _save_checkpoint(self, context: TrainContext):
                     **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 @@ def _configure_rollout(self, context: TrainContext, strategy_kwargs: dict) -> No
                 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 @@ def _trainer_loop(self, param_path: Optional[str] = None, resume: bool = False):
 
                         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: