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:
@@ -231,6 +231,12 @@ class CheckpointCallback(TrainCallback):
|
||||
obj = getattr(context, name, None)
|
||||
if obj:
|
||||
extra[name] = obj.state_dict()
|
||||
critic = getattr(context.strategy, "critic", None)
|
||||
if critic is not None:
|
||||
extra["value_model"] = critic.state_dict()
|
||||
critic_optimizer = getattr(context.strategy, "critic_optimizer", None)
|
||||
if critic_optimizer is not None:
|
||||
extra["value_optimizer"] = critic_optimizer.state_dict()
|
||||
return extra
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user