perf: reuse rollout behavior logprobs

astrai/trainer/strategy.py

@@ -535,12 +535,12 @@ class GRPOStrategy(BaseStrategy):
     broadcast across all response tokens.  The loss is computed **only on
     response tokens** — prompt tokens are masked out.
 
-    Three model roles are distinguished:
+    Three policy roles are distinguished:
 
     * **Policy** ``self.model`` — the model being trained.
-    * **Old policy** ``self.old_model`` — the behaviour policy that generated
-      the responses.  Used for the importance sampling ratio
-      ``ρ = π_θ / π_old``.  Synced externally after each data-generation round.
+    * **Behaviour policy** — represented by per-token ``logprobs_old`` captured
+      during online rollout.  Offline batches may instead use ``self.old_model``
+      as a compatibility fallback.
     * **Reference model** ``self.ref_model`` — a frozen copy of the initial
       policy (typically the SFT checkpoint) used **only** for the KL
       regularisation term.  It is never updated during training.
@@ -550,7 +550,7 @@ def __init__(
         self,
         model: nn.Module,
         device: str,
-        old_model: nn.Module,
+        old_model: Optional[nn.Module],
         ref_model: nn.Module,
         clip_eps: float = 0.2,
         kl_coef: float = 0.01,
@@ -566,6 +566,8 @@ def __init__(
 
     def sync_old_model(self):
         """Copy current policy weights to old model."""
+        if self.old_model is None:
+            raise RuntimeError("Cannot sync an unconfigured old policy model")
         state_dict = self.executor.unwrap_model(self.model)
         if self.executor.use_distributed:
             state_dict = broadcast_state_dict(state_dict)
@@ -580,6 +582,22 @@ def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput:
         rewards = batch["rewards"]
 
         batch_size, group_size, response_len = responses.shape
+        behavior_logprobs = batch.get("logprobs_old")
+        if behavior_logprobs is not None:
+            if behavior_logprobs.shape != responses.shape:
+                raise ValueError(
+                    "logprobs_old shape must match responses: "
+                    f"got {tuple(behavior_logprobs.shape)}, "
+                    f"expected {tuple(responses.shape)}"
+                )
+            if not torch.isfinite(behavior_logprobs).all():
+                raise ValueError("logprobs_old must contain only finite values")
+            behavior_logprobs = behavior_logprobs.detach().float()
+        elif self.old_model is None:
+            raise ValueError(
+                "GRPO batches must provide logprobs_old when no old_model is configured"
+            )
+
         responses_flat = responses.view(-1, response_len)
         masks_flat = masks.view(-1, response_len)
         prompt_expanded = prompts.unsqueeze(1).repeat(1, group_size, 1).flatten(0, 1)
@@ -620,11 +638,14 @@ def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput:
         aux_loss = policy_output["aux_loss"]
         token_log_probs_policy = token_log_probs_policy[:, prompt_len - 1 :]
         with torch.no_grad():
-            old_output = get_logprobs(
-                self.old_model, full_sequences, attn_mask, full_masks, "none"
-            )
-            token_log_probs_old = old_output["logprobs"]
-            token_log_probs_old = token_log_probs_old[:, prompt_len - 1 :]
+            if behavior_logprobs is None:
+                old_output = get_logprobs(
+                    self.old_model, full_sequences, attn_mask, full_masks, "none"
+                )
+                token_log_probs_old = old_output["logprobs"]
+                token_log_probs_old = token_log_probs_old[:, prompt_len - 1 :]
+            else:
+                token_log_probs_old = behavior_logprobs
             ref_output = get_logprobs(
                 self.ref_model, full_sequences, attn_mask, full_masks, "none"
             )
@@ -680,12 +701,9 @@ def prepare_from_rollout(self, result: RolloutResult) -> Dict[str, Tensor]:
             "responses": result.responses,
             "masks": result.response_mask,
             "rewards": result.rewards,
+            "logprobs_old": result.logprobs_old,
         }
 
-    def _on_rollout_refresh(self):
-        """Sync the behaviour policy whenever a fresh rollout arrives."""
-        self.sync_old_model()
-
 
 # Factory aliases: online variants use the same strategy class; the
 # ``RolloutRunner`` is injected by ``TrainContextBuilder`` to enable

astrai/trainer/train_context.py

@@ -287,13 +287,15 @@ def _create_strategy(self, context: TrainContext, executor: BaseExecutor) -> dic
                 model=context.model,
                 device=get_current_device(),
             )
-        if cfg.strategy in ("grpo", "online_grpo"):
+        if cfg.strategy == "grpo":
             kwargs["old_model"] = create_ref_model(
                 cfg.model_fn,
                 executor=executor,
                 model=context.model,
                 device=get_current_device(),
             )
+        elif cfg.strategy == "online_grpo":
+            kwargs["old_model"] = None
         context.strategy = StrategyFactory.create(
             cfg.strategy,
             model=context.model,

docs/developer/internals.md

@@ -84,7 +84,7 @@ $$ \text{Advantage}_i = \frac{r_i - \mu}{\sigma + \epsilon} $$
 
 $$ L_{\text{GRPO}} = -\mathbb{E}_t\left[\min\left(\rho_t A,\; \text{clip}\left(\rho_t, 1-\epsilon, 1+\epsilon\right)A\right)\right] + \lambda \cdot \mathbb{E}_t\left[\frac{\pi_{\text{ref}}}{\pi_\theta} - \log\frac{\pi_{\text{ref}}}{\pi_\theta} - 1\right] $$
 
-Where $\rho_t = \pi_\theta(a_t|s_t) / \pi_{\text{old}}(a_t|s_t)$ is the per-token importance sampling ratio. Advantages are derived from scalar per-response rewards, group-normalized, and broadcast across all response tokens. Only response tokens contribute to the loss.
+Where $\rho_t = \pi_\theta(a_t|s_t) / \pi_{\text{old}}(a_t|s_t)$ is the per-token importance sampling ratio. Online rollout records $\log \pi_{\text{old}}$ when each token is sampled and reuses those values directly during training; offline batches may fall back to a synchronized `old_model`. Advantages are derived from scalar per-response rewards, group-normalized, and broadcast across all response tokens. Only response tokens contribute to the loss.
 
 Parameters: `group_size=4`, `clip_eps=0.2`, `kl_coef=0.01`.
 

docs/guides/training.md

@@ -148,23 +148,28 @@ $$
 
 where $\rho_t = \pi_\theta(a_t|s_t) / \pi_{\text{old}}(a_t|s_t)$ is the
 per-token importance sampling ratio against the behaviour policy
-(`old_model`, synced externally between data-generation rounds) and the
-expectations are over valid response tokens. The KL term regularises
-$\pi_\theta$ towards a frozen reference model (`ref_model`, typically
-the SFT checkpoint).
+and the expectations are over valid response tokens. Online GRPO reuses the
+per-token `logprobs_old` captured by the rollout sampler, avoiding an
+`old_model` copy and a repeated forward pass. Offline GRPO keeps `old_model` as
+a compatibility fallback. The KL term regularises $\pi_\theta$ towards a frozen
+reference model (`ref_model`, typically the SFT checkpoint).
 
-Parameters: `group_size=4`, `clip_eps=0.2`, `kl_coef=0.01`. External sync of `old_model` weights via `sync_old_model()` between data-generation rounds.
+Parameters: `group_size=4`, `clip_eps=0.2`, `kl_coef=0.01`. Offline callers that
+do not provide `logprobs_old` must sync `old_model` weights via
+`sync_old_model()` between data-generation rounds.
 
-Keys: `prompts`, `responses`, `masks`, `rewards`.
+Keys: `prompts`, `responses`, `masks`, `rewards`, and optional
+`logprobs_old` (required when `old_model` is not configured).
 
 ### Online Rollout
 
 `online_grpo` and `online_dpo` use the respective GRPO and DPO strategies with
 a `RolloutRunner`. The runner renders prompts through the tokenizer chat
 template, generates grouped responses through `InferenceScheduler`, then scores
 them with a `BaseRewardModel`. It refreshes cached rollouts every
-`rollout_interval` optimizer steps. `online_grpo` synchronizes `old_model` when
-a fresh rollout is produced.
+`rollout_interval` optimizer steps. `online_grpo` carries the sampler's aligned
+behaviour log-probabilities into the loss, so it does not allocate or synchronize
+a separate old-policy model.
 
 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

tests/trainer/test_grpo_strategy.py

@@ -71,6 +71,44 @@ def test_grpo_loss_backward(grpo_strategy):
     assert has_grad
 
 
+def test_grpo_reuses_supplied_behavior_logprobs(grpo_strategy):
+    """A rollout batch must not forward the old policy again."""
+    strategy, device = grpo_strategy
+
+    class _FailingOldPolicy(torch.nn.Module):
+        def forward(self, *args, **kwargs):
+            raise AssertionError("old policy forward should not run")
+
+    strategy.old_model = _FailingOldPolicy()
+    batch = _make_batch(device=device)
+    batch["logprobs_old"] = torch.zeros_like(batch["responses"], dtype=torch.float)
+
+    loss = strategy.compute_loss(batch)
+    assert torch.isfinite(loss).item()
+
+
+def test_grpo_requires_behavior_source(grpo_strategy):
+    strategy, device = grpo_strategy
+    strategy.old_model = None
+    with pytest.raises(ValueError, match="must provide logprobs_old"):
+        strategy.compute_loss(_make_batch(device=device))
+
+
+@pytest.mark.parametrize("invalid", ["shape", "nonfinite"])
+def test_grpo_rejects_invalid_behavior_logprobs(grpo_strategy, invalid):
+    strategy, device = grpo_strategy
+    batch = _make_batch(device=device)
+    if invalid == "shape":
+        batch["logprobs_old"] = torch.zeros(1, device=device)
+        match = "shape must match responses"
+    else:
+        batch["logprobs_old"] = torch.zeros_like(batch["responses"], dtype=torch.float)
+        batch["logprobs_old"][0, 0, 0] = float("nan")
+        match = "only finite values"
+    with pytest.raises(ValueError, match=match):
+        strategy.compute_loss(batch)
+
+
 @pytest.mark.parametrize("model_name", ["ref_model", "old_model"])
 def test_grpo_frozen_models_not_updated(grpo_strategy, model_name):
     """Backward should not populate gradients on ref_model or old_model."""

tests/trainer/test_online_e2e.py

@@ -7,6 +7,7 @@
 import torch
 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.trainer.rollout import BaseRewardModel
@@ -87,8 +88,19 @@ 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):
+def test_online_rollout_end_to_end(
+    base_test_env, strategy, strategy_kwargs, monkeypatch
+):
     """Run one epoch of online RL rollout with KV-cache-backed generation."""
+    created_reference_models = []
+    create_ref_model = train_context.create_ref_model
+
+    def track_reference_model(*args, **kwargs):
+        created_reference_models.append(strategy)
+        return create_ref_model(*args, **kwargs)
+
+    monkeypatch.setattr(train_context, "create_ref_model", track_reference_model)
+
     test_dir = base_test_env["test_dir"]
     device = base_test_env["device"]
     tokenizer = base_test_env["tokenizer"]
@@ -126,3 +138,4 @@ def test_online_rollout_end_to_end(base_test_env, strategy, strategy_kwargs):
     trainer.train(param_path=test_dir)
 
     assert os.path.isdir(os.path.join(test_dir, "ckpt"))
+    assert len(created_reference_models) == 1

tests/trainer/test_online_strategy.py

@@ -60,12 +60,11 @@ def swap_result(self, result):
 
 def _make_grpo(device, executor=None):
     model, _ = make_model(device)
-    old_model = make_frozen(model, device)
     ref_model = make_frozen(model, device)
     return GRPOStrategy(
         model=model,
         device=device,
-        old_model=old_model,
+        old_model=None,
         ref_model=ref_model,
         clip_eps=0.2,
         kl_coef=0.01,
@@ -134,6 +133,7 @@ def test_grpo_prepare_from_rollout_mapping(device):
     assert batch["responses"] is r.responses
     assert batch["masks"] is r.response_mask
     assert batch["rewards"] is r.rewards
+    assert batch["logprobs_old"] is r.logprobs_old
 
 
 def test_dpo_prepare_from_rollout_conditions_responses_on_prompt(device):
@@ -190,13 +190,14 @@ def test_dpo_prepare_from_rollout_same_response_keeps_distinct_prompts():
     assert not batch["rejected_mask"][:, :3].any()
 
 
-def test_call_without_runner_falls_back_to_compute_loss_grpo(device):
+def test_call_without_runner_accepts_behavior_logprobs_grpo(device):
     strat = _make_grpo(device)
     batch = {
         "prompts": torch.randint(3, 200, (2, 4), device=device),
         "responses": torch.randint(3, 200, (2, 4, 6), device=device),
         "masks": torch.ones(2, 4, 6, device=device),
         "rewards": torch.randn(2, 4, device=device),
+        "logprobs_old": torch.zeros(2, 4, 6, device=device),
     }
     loss = strat(batch)["loss"]
     assert torch.isfinite(loss).item()
@@ -225,25 +226,19 @@ def test_call_invokes_runner_each_time(device):
     assert runner.calls == 2
 
 
-def test_grpo_syncs_old_model_on_first_rollout(device):
+def test_grpo_reuses_rollout_logprobs_without_old_model(device):
     strat = _make_grpo(device)
-    runner = _RecordingRunner(_make_rollout_result(device=device))
+    result = _make_rollout_result(device=device)
+    result.logprobs_old.normal_().requires_grad_()
+    runner = _RecordingRunner(result)
     strat.set_rollout_runner(runner)
-    with torch.no_grad():
-        for p in strat.model.parameters():
-            p.add_(0.1)
-    old_before = {k: v.clone() for k, v in strat.old_model.state_dict().items()}
-    strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})
-    old_after = strat.old_model.state_dict()
-    synced = any(
-        not torch.allclose(old_before[k], old_after[k])
-        for k in old_before
-        if k in old_after
-    )
-    assert synced
+    assert strat.old_model is None
+    loss = strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})["loss"]
+    loss.backward()
+    assert result.logprobs_old.grad is None
 
 
-def test_grpo_no_resync_when_same_cached_result(device):
+def test_grpo_reuses_same_cached_result(device):
     strat = _make_grpo(device)
     runner = _RecordingRunner(_make_rollout_result(device=device))
     strat.set_rollout_runner(runner)
@@ -255,7 +250,7 @@ def test_grpo_no_resync_when_same_cached_result(device):
     assert runner.step_calls == 2
 
 
-def test_grpo_resync_when_new_rollout_result(device):
+def test_grpo_accepts_new_rollout_result(device):
     strat = _make_grpo(device)
     runner = _RecordingRunner(_make_rollout_result(device=device))
     strat.set_rollout_runner(runner)