perf: reuse rollout behavior logprobs
Feed sampler-aligned behavior log-probabilities directly into online GRPO instead of allocating, synchronizing, and forwarding a duplicate old-policy model. Keep the old-model path as an offline compatibility fallback and validate supplied rollout tensors before loss computation.
This commit is contained in:
0z5a
ViperEkura
+32
-14
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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`.
|
||||
|
||||
|
||||
+13
-8
@@ -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
|
||||
|
||||
@@ -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."""
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user