diff --git a/astrai/trainer/strategy.py b/astrai/trainer/strategy.py index c0bb8ba..2321c7d 100644 --- a/astrai/trainer/strategy.py +++ b/astrai/trainer/strategy.py @@ -542,12 +542,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. @@ -557,7 +557,7 @@ class GRPOStrategy(BaseStrategy): 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, @@ -573,6 +573,8 @@ class GRPOStrategy(BaseStrategy): 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) @@ -587,6 +589,22 @@ class GRPOStrategy(BaseStrategy): 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) @@ -627,11 +645,14 @@ class GRPOStrategy(BaseStrategy): 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" ) @@ -687,12 +708,9 @@ class GRPOStrategy(BaseStrategy): "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 diff --git a/astrai/trainer/train_context.py b/astrai/trainer/train_context.py index 3b59dc7..5753c92 100644 --- a/astrai/trainer/train_context.py +++ b/astrai/trainer/train_context.py @@ -287,13 +287,15 @@ class TrainContextBuilder: 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, diff --git a/docs/developer/internals.md b/docs/developer/internals.md index 71379dd..ccfdd8f 100644 --- a/docs/developer/internals.md +++ b/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`. diff --git a/docs/guides/training.md b/docs/guides/training.md index d6475a2..0d267f1 100644 --- a/docs/guides/training.md +++ b/docs/guides/training.md @@ -148,14 +148,18 @@ $$ 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 @@ -163,8 +167,9 @@ Keys: `prompts`, `responses`, `masks`, `rewards`. 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. Every successful optimizer step advances a monotonic `policy_version` and acknowledges the shared-model weight update to the rollout scheduler. The diff --git a/tests/trainer/test_grpo_strategy.py b/tests/trainer/test_grpo_strategy.py index d7c531b..473871b 100644 --- a/tests/trainer/test_grpo_strategy.py +++ b/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.""" diff --git a/tests/trainer/test_online_e2e.py b/tests/trainer/test_online_e2e.py index 9f1eac8..df8b817 100644 --- a/tests/trainer/test_online_e2e.py +++ b/tests/trainer/test_online_e2e.py @@ -7,6 +7,7 @@ import pytest 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 @@ _ONLINE_STRATEGIES = [ @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 diff --git a/tests/trainer/test_online_strategy.py b/tests/trainer/test_online_strategy.py index ad4732e..e66a500 100644 --- a/tests/trainer/test_online_strategy.py +++ b/tests/trainer/test_online_strategy.py @@ -67,12 +67,11 @@ class _RecordingRunner: 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, @@ -141,6 +140,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): @@ -197,13 +197,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() @@ -232,25 +233,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) @@ -262,7 +257,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)