feat: add online ppo with value-model critic and gae advantages

- register online_ppo train type backed by PPOStrategy: token-level clipped surrogate over GAE advantages plus masked value regression against rollout-pinned returns, with explained-variance metrics
- fold the reference-KL penalty (k3 estimator) into per-token rewards before GAE and pin advantages/returns on RolloutResult so replayed gradient steps optimize fixed targets
- add self-contained ValueModel critic with a zero-initialized value head and backbone warm-started from policy weights; AutoRegressiveLM stays untouched and trunk parity is pinned by tests
- step the critic's own optimizer outside the policy-version lock with the same max_grad_norm clipping as the policy
- persist critic state as value_model.pt/value_optimizer.pt checkpoint extras; resume restores it, fails loudly when missing, and the train.sh completeness check requires the extras for online_ppo configs
- extract shared rollout sequence/logprob helpers from GRPO (behavior unchanged) and add ppo_gamma/ppo_gae_lambda/ppo_vf_coef CLI options
This commit is contained in:
2026-09-05 01:59:50 +08:00
parent 816b96a58a
commit 350e4a1849
17 changed files with 1390 additions and 72 deletions
+12
View File
@@ -170,6 +170,18 @@ them with a `BaseRewardModel`. It refreshes cached rollouts every
behaviour log-probabilities into the loss, so it does not allocate or synchronize
a separate old-policy model.
`online_ppo` is actor-critic PPO on the same rollout pipeline. A `ValueModel`
critic (backbone warm-started from the policy, zero-initialized value head)
scores the rollout states; advantages come from GAE(`--ppo_gamma`,
`--ppo_gae_lambda`) with the terminal reward on each response's last token and
the reference-KL penalty (k3 estimator, `--grpo_kl_coef`) folded into per-token
rewards. Advantages and returns are computed once per rollout and pinned on the
`RolloutResult`, so replayed steps optimize fixed targets. The critic has its
own optimizer, stepped outside the policy-version lock, and persists as
`value_model.pt`/`value_optimizer.pt` checkpoint extras — resume without them
fails loudly, and `scripts/train.sh` treats a PPO checkpoint as incomplete when
they are missing.
Every successful optimizer step mutates the shared model and advances its
monotonic `policy_version` under the same generation lock. The scheduler
invalidates reusable KV prefixes before accepting the new version, so an async