refactor: separate old policy and ref model in GRPO strategy

- Split single ref_model into old_model (importance sampling ratio) and ref_model (frozen KL regularizer)
- Move ref_model/old_model creation from strategy __init__ to TrainContextBuilder, pass as explicit parameters
- Remove periodic sync_ref_model + sync_interval; add sync_old_model for external rollout loop to call
- DPOStrategy also receives ref_model from builder
- Fix std to use unbiased=False (population std per GRPO paper)
- Remove redundant tests (test_grpo_kl_zero_at_init, test_grpo_no_sync_interval_param)
- Remove --grpo_sync_interval CLI arg

astrai/trainer/strategy.py

@@ -223,14 +223,13 @@ class DPOStrategy(BaseStrategy):
         self,
         model: nn.Module,
         device: str,
+        ref_model: nn.Module,
         beta: float = 0.1,
         reduction: str = "mean",
         **kwargs,
     ):
         super().__init__(model, device, **kwargs)
-        self.ref_model = create_ref_model(
+        self.ref_model = ref_model
-            self.model_fn, self.executor.unwrap_model(model)
-        ).to(device=self.device)
         self.beta = beta
         self.reduction = reduction
 
@@ -272,40 +271,40 @@ class GRPOStrategy(BaseStrategy):
     broadcast across all response tokens.  The loss is computed **only on
     response tokens** — prompt tokens are masked out.
 
-    The strategy expects offline-collected batches (``responses`` / ``rewards``
+    Three model roles are distinguished:
-    pre-generated by the current or a recent policy).  Call ``sync_ref_model()``
+
-    after each data-generation round so ``ref_model`` tracks the sampling policy.
+    * **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.
+    * **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.
     """
 
     def __init__(
         self,
         model: nn.Module,
         device: str,
+        old_model: nn.Module,
+        ref_model: nn.Module,
         clip_eps: float = 0.2,
         kl_coef: float = 0.01,
         group_size: int = 4,
-        sync_interval: int = 200,
         **kwargs,
     ):
         super().__init__(model, device, **kwargs)
-        self.ref_model = create_ref_model(
+        self.old_model = old_model
-            self.model_fn, self.executor.unwrap_model(model)
+        self.ref_model = ref_model
-        ).to(device=self.device)
         self.clip_eps = clip_eps
         self.kl_coef = kl_coef
         self.group_size = group_size
-        self.sync_interval = sync_interval
-        self._step = 0
 
-    def sync_ref_model(self):
+    def sync_old_model(self):
-        """Copy current model weights to ref model."""
+        """Copy current policy weights to old model."""
-        self.ref_model.load_state_dict(self.executor.unwrap_model(self.model))
+        self.old_model.load_state_dict(self.executor.unwrap_model(self.model))
 
     def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
-        self._step += 1
-        if self._step % self.sync_interval == 0:
-            self.sync_ref_model()
-
         batch = move_to_device(batch, self.device)
         prompts = batch["prompts"]
         responses = batch["responses"]
@@ -332,25 +331,29 @@ class GRPOStrategy(BaseStrategy):
             self.model, full_sequences, full_masks, "none"
         )[:, prompt_len - 1 :]
         with torch.no_grad():
+            token_log_probs_old = get_logprobs(
+                self.old_model, full_sequences, full_masks, "none"
+            )[:, prompt_len - 1 :]
             token_log_probs_ref = get_logprobs(
                 self.ref_model, full_sequences, full_masks, "none"
             )[:, prompt_len - 1 :]
 
         # Reshape to [B, G, response_len]
         token_log_probs_policy = token_log_probs_policy.view(batch_size, group_size, -1)
+        token_log_probs_old = token_log_probs_old.view(batch_size, group_size, -1)
         token_log_probs_ref = token_log_probs_ref.view(batch_size, group_size, -1)
         token_masks = masks_flat.view(batch_size, group_size, -1).float()
 
         # Group-normalized advantages from scalar per-response rewards.
         eps = 1e-8
         mean = rewards.mean(dim=-1, keepdim=True)
-        std = rewards.std(dim=-1, keepdim=True)
+        std = rewards.std(dim=-1, keepdim=True, unbiased=False)
         advantages = (rewards - mean) / (std + eps)
         # Broadcast scalar advantage to every response token: [B, G, 1]
         advantages = advantages.unsqueeze(-1)
 
-        # Token-level ratio and PPO clipping.
+        # Token-level ratio (π_θ / π_old) and PPO clipping.
-        log_ratio = token_log_probs_policy - token_log_probs_ref
+        log_ratio = token_log_probs_policy - token_log_probs_old
         ratio = torch.exp(log_ratio)
 
         surr1 = ratio * advantages
@@ -359,8 +362,10 @@ class GRPOStrategy(BaseStrategy):
         token_count = token_masks.sum().clamp(min=1.0)
         policy_loss = (per_token_policy_loss * token_masks).sum() / token_count
 
-        # KL penalty with k1 estimator (non-negative): r - log(r) - 1, r=π_ref/π_θ.
+        # KL penalty to frozen reference model with k1 estimator (non-negative):
-        r = torch.exp(-log_ratio)
+        # k1 = π_ref / π_θ - log(π_ref / π_θ) - 1, where π_ref / π_θ = exp(log_ref - log_policy).
+        log_ref_ratio = token_log_probs_ref - token_log_probs_policy
+        r = torch.exp(log_ref_ratio)
         kl_per_token = r - torch.log(r + eps) - 1.0
         kl_penalty = self.kl_coef * (kl_per_token * token_masks).sum() / token_count
 

astrai/trainer/train_context.py

@@ -13,7 +13,7 @@ from astrai.parallel.executor import BaseExecutor, ExecutorFactory
 from astrai.parallel.setup import get_current_device, get_rank, get_world_size
 from astrai.protocols import OptimizerProtocol, SchedulerProtocol
 from astrai.serialization import Checkpoint, load_json
-from astrai.trainer.strategy import BaseStrategy, StrategyFactory
+from astrai.trainer.strategy import BaseStrategy, StrategyFactory, create_ref_model
 
 
 @dataclass
@@ -177,13 +177,27 @@ class TrainContextBuilder:
                     if obj is not None:
                         obj.load_state_dict(extra[name])
 
+        strategy_kwargs = dict(cfg.extra_kwargs)
+
+        if cfg.strategy in ("dpo", "grpo"):
+            ref_model = create_ref_model(
+                cfg.model_fn, executor.unwrap_model(context.model)
+            ).to(device=device)
+            strategy_kwargs["ref_model"] = ref_model
+
+        if cfg.strategy == "grpo":
+            old_model = create_ref_model(
+                cfg.model_fn, executor.unwrap_model(context.model)
+            ).to(device=device)
+            strategy_kwargs["old_model"] = old_model
+
         context.strategy = StrategyFactory.create(
             cfg.strategy,
             model=context.model,
             device=device,
             executor=executor,
             model_fn=cfg.model_fn,
-            **cfg.extra_kwargs,
+            **strategy_kwargs,
         )
 
         return context

scripts/tools/train.py

@@ -252,12 +252,6 @@ def parse_args() -> argparse.Namespace:
         default="checkpoint/logs",
         help="Directory for metric logs.",
     )
-    parser.add_argument(
-        "--grpo_sync_interval",
-        type=int,
-        default=200,
-        help="GRPO ref model sync interval (steps).",
-    )
     parser.add_argument(
         "--start_epoch", type=int, default=0, help="Start epoch for training."
     )
@@ -444,7 +438,6 @@ def train(
         "clip_eps": kwargs.pop("grpo_clip_eps"),
         "kl_coef": kwargs.pop("grpo_kl_coef"),
         "group_size": kwargs.pop("group_size"),
-        "sync_interval": kwargs.pop("grpo_sync_interval"),
     }
 
     executor_kwargs = {

tests/trainer/test_grpo_strategy.py

@@ -56,19 +56,32 @@ def _make_batch(
     }
 
 
+def _make_frozen_copy(model, device):
+    """Create a frozen copy of ``model`` with independent weights loaded."""
+    config = _make_config()
+    copy = AutoRegressiveLM(config).to(device=device)
+    copy.load_state_dict(model.state_dict())
+    copy.requires_grad_(False)
+    copy.eval()
+    return copy
+
+
 @pytest.fixture
 def grpo_strategy():
     """Build a GRPOStrategy with a small real model and fake executor."""
     device = "cuda" if torch.cuda.is_available() else "cpu"
     model, config = _make_model(device)
+    old_model = _make_frozen_copy(model, device)
+    ref_model = _make_frozen_copy(model, device)
 
     strategy = GRPOStrategy(
         model=model,
         device=device,
+        old_model=old_model,
+        ref_model=ref_model,
         clip_eps=0.2,
         kl_coef=0.01,
         group_size=4,
-        sync_interval=200,
         model_fn=lambda c=config: AutoRegressiveLM(c).to(device=device),
         executor=_FakeExecutor(),
     )
@@ -108,6 +121,16 @@ def test_grpo_ref_model_not_updated(grpo_strategy):
         assert p.grad is None
 
 
+def test_grpo_old_model_not_updated(grpo_strategy):
+    """Backward should not populate gradients on old_model."""
+    strategy, device = grpo_strategy
+    batch = _make_batch(device=device)
+    loss = strategy.compute_loss(batch)
+    loss.backward()
+    for p in strategy.old_model.parameters():
+        assert p.grad is None
+
+
 def test_grpo_prompt_tokens_masked(grpo_strategy):
     """When only prompt-equivalent tokens are unmasked (response mask all 0),
     the policy loss should be zero (no valid tokens contribute)."""
@@ -127,45 +150,32 @@ def test_grpo_identical_rewards_zero_advantage(grpo_strategy):
     batch = _make_batch(device=device)
     batch["rewards"] = torch.ones(batch["rewards"].shape, device=device)
     loss = strategy.compute_loss(batch)
-    # At init policy == ref, so ratio == 1, KL == 0; advantage == 0.
+    # At init policy == old == ref, so ratio == 1, KL == 0; advantage == 0.
     assert loss.item() == pytest.approx(0.0, abs=1e-5)
 
 
-def test_grpo_kl_zero_at_init(grpo_strategy):
+def test_grpo_sync_old_model(grpo_strategy):
-    """At initialization policy == ref_model, so KL penalty must be 0."""
+    """sync_old_model copies current policy weights into old_model."""
     strategy, device = grpo_strategy
-    batch = _make_batch(device=device)
+    # Perturb policy model so it differs from old.
-    # Make rewards distinct so advantages are non-zero (isolates KL term).
-    loss = strategy.compute_loss(batch)
-    # KL term is 0 at init; loss is purely policy surrogate.
-    # With ratio==1, surr1==surr2==advantage, so policy_loss = -mean(|adv|).
-    # Just assert KL portion is negligible by checking loss is finite and
-    # re-running after a model update increases loss magnitude.
-    assert torch.isfinite(loss).item()
-
-
-def test_grpo_sync_ref_model(grpo_strategy):
-    """sync_ref_model copies current policy weights into ref_model."""
-    strategy, device = grpo_strategy
-    # Perturb policy model so it differs from ref.
     with torch.no_grad():
         for p in strategy.model.parameters():
             p.add_(0.05)
-    # ref_model should still hold original weights (differ from policy).
+    # old_model should still hold original weights (differ from policy).
     policy_sd = strategy.model.state_dict()
-    ref_sd = strategy.ref_model.state_dict()
+    old_sd = strategy.old_model.state_dict()
     differs_before = any(
-        not torch.allclose(policy_sd[k], ref_sd[k]) for k in policy_sd if k in ref_sd
+        not torch.allclose(policy_sd[k], old_sd[k]) for k in policy_sd if k in old_sd
     )
     assert differs_before
 
-    strategy.sync_ref_model()
+    strategy.sync_old_model()
 
-    ref_sd_after = strategy.ref_model.state_dict()
+    old_sd_after = strategy.old_model.state_dict()
     matches = all(
-        torch.allclose(policy_sd[k], ref_sd_after[k])
+        torch.allclose(policy_sd[k], old_sd_after[k])
         for k in policy_sd
-        if k in ref_sd_after
+        if k in old_sd_after
     )
     assert matches