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:
+11
-5
@@ -14,7 +14,7 @@
|
||||
| Parameter | Description | Default |
|
||||
|-----------|-------------|---------|
|
||||
| `--config`, `-c` | YAML config file; explicit CLI options override YAML values | None |
|
||||
| `--train_type` | Training type (`seq`, `sft`, `dpo`, `grpo`, `online_grpo`, `online_dpo`) | required |
|
||||
| `--train_type` | Training type (`seq`, `sft`, `dpo`, `grpo`, `online_grpo`, `online_dpo`, `online_ppo`) | required |
|
||||
| `--data_root_path` | Dataset root directory | required |
|
||||
| `--param_path` | Model parameters or checkpoint path | required |
|
||||
| `--resume` | Resume training from `--param_path` | False |
|
||||
@@ -139,16 +139,22 @@ with `--optimizer=muon_adamw`.
|
||||
|-----------|-------------|---------|---------|
|
||||
| `--dpo_beta` | DPO beta value | 0.1 | `dpo`, `online_dpo` |
|
||||
| `--label_smoothing` | Label smoothing for cross-entropy loss | 0.0 | `seq`, `sft` |
|
||||
| `--group_size` | GRPO/rollout group size | 4 | `grpo`, `online_grpo`, `online_dpo` |
|
||||
| `--grpo_clip_eps` | GRPO clipping epsilon | 0.2 | `grpo`, `online_grpo` |
|
||||
| `--grpo_kl_coef` | GRPO KL penalty coefficient | 0.01 | `grpo`, `online_grpo` |
|
||||
| `--group_size` | GRPO/rollout group size | 4 | `grpo`, `online_grpo`, `online_dpo`, `online_ppo` |
|
||||
| `--grpo_clip_eps` | Clipping epsilon for the PPO-style surrogate loss | 0.2 | `grpo`, `online_grpo`, `online_ppo` |
|
||||
| `--grpo_kl_coef` | KL penalty coefficient | 0.01 | `grpo`, `online_grpo`, `online_ppo` |
|
||||
| `--ppo_gamma` | PPO reward discount factor | 1.0 | `online_ppo` |
|
||||
| `--ppo_gae_lambda` | PPO GAE bias/variance trade-off | 0.95 | `online_ppo` |
|
||||
| `--ppo_vf_coef` | PPO value-loss coefficient | 0.5 | `online_ppo` |
|
||||
| `--neftune_alpha` | NEFTune noise alpha (0=disabled, typical: 5.0) | 0.0 | `sft` |
|
||||
|
||||
### Online Rollout
|
||||
|
||||
`online_grpo` and `online_dpo` are factory aliases for the existing `grpo` and
|
||||
`dpo` strategy classes; online behavior is enabled by rollout components rather
|
||||
than separate strategy subclasses. These options apply to the online aliases.
|
||||
than separate strategy subclasses. `online_ppo` is a dedicated actor-critic
|
||||
strategy: a `ValueModel` critic supplies GAE advantages, and its state persists
|
||||
as `value_model.pt`/`value_optimizer.pt` checkpoint extras (required for
|
||||
resume). These options apply to the online strategies.
|
||||
Online strategies require
|
||||
a `BaseRewardModel` factory in `TrainConfig`; `train.py` does not currently
|
||||
provide a command-line option for configuring one.
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user