feat: add standard PPO training with GAE and value critic (Transformers/FSDP)

[xxyyrr598]

PR type

• New Feature
• Bug Fix
• Document Updates
• More Models or Datasets Support

PR information

Summary

This PR adds the first standard PPO + GAE training pipeline to Twinkle, on top of the existing GRPO/DPO RL stack. It introduces the missing PPO components — a token-level GAE advantage estimator, a full-parameter value critic, and a clipped value loss — and provides an end-to-end GSM8K cookbook with LoRA policy + full-parameter critic + vLLM sampler.

What's included

New components

• cookbook/rl/ppo/ — end-to-end GSM8K PPO training script & shell launcher:
  • LoRA policy (FSDP) + full-parameter critic (FSDP) + vLLM sampler on separate GPU groups
  • Frozen base model as reference policy via forward_only(disable_lora=True) (no extra ref GPUs)
  • ppo_epochs mini-batch updates over fixed rollouts
• src/twinkle/advantage/gae.py — GAEAdvantage, token-level GAE with:
  • terminal bootstrap = 0, padding masked out
  • per-token KL shaping reward (logp_old - logp_ref) with kl_coef
  • optional advantage normalization
• src/twinkle/loss/value.py — PPOValueLoss, clipped value objective max((v-r)², (clip(v,v_old)-r)²)
• src/twinkle/model/transformers/value_model.py — TransformersValueModel, causal-LM backbone with a hidden_size -> 1 value head replacing the LM head
• src/twinkle/metric/ppo.py — PPOValueMetric (value/return mean, value clip ratio, explained variance)
• PPOLoss(GRPOLoss) / PPOMetric(GRPOMetric) — semantic aliases reusing the existing clipped policy objective & ratio/clip statistics

Framework generalization (backward-compatible)

• ModelOutput gains a values field; forward/forward_only can emit values when the loss declares require_values
• processor packing, nccl_safe wrapper, and loss/base handle the new values tensor
• New CLI args: --critic-model-gpus, --ppo-epochs, --gamma, --gae-lambda, --kl-coef, --normalize-advantages, --critic-learning-rate, --value-clip

Tests

• tests/advantage/test_gae.py — single/multi-step GAE, terminal bootstrap, padding mask, normalization, sparse reward + KL shaping
• tests/loss/test_ppo.py — PPOLoss parity with GRPOLoss, clipped/unclipped value loss, ragged token alignment
• tests/model/test_value_model.py — value-head forward/backward, save/load
• tests/cli/test_cli.py — new PPO CLI args

Docs

• README / README_ZH cookbook table updated with the PPO entry

Usage

cd cookbook/rl/ppo
bash ppo.sh   # default 12 GPUs: 4 policy + 4 critic + 4 sampler

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.