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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user