feat: add standard PPO training with GAE and value critic (Transformers/FSDP) - #256
feat: add standard PPO training with GAE and value critic (Transformers/FSDP)#256xxyyrr598 wants to merge 7 commits into
Conversation
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
| del logits | ||
| if loss_require_values: | ||
| values = outputs['logits'] | ||
| outputs['values'] = values.squeeze(-1) if values.shape[-1] == 1 else values |
There was a problem hiding this comment.
此处无需对 values 额外 shift。Twinkle 的 labels 已遵循 causal-LM 的预移位约定，logits[t] 直接与 labels[t] 对齐；Value Model 的 values[t] 也表示预测该 token action 前状态的 V(s_t)。PPO 中 logprob、value、reward 均按同一 labels mask 选取，GAE 所需的 V(s_{t+1}) 在 advantage 计算阶段通过 next value 获取，末位置 bootstrap 为 0。若在此处单独 shift values，会破坏该对齐关系。
| # Auto-detect device resources for accelerators that Ray doesn't | ||
| # natively recognise (e.g. Ascend NPU). Ray auto-detects GPU via | ||
| # nvidia-smi, but has no built-in detection for NPU / other devices, | ||
| # so without explicit resources Ray nodes report 0 NPU and the | ||
| # subsequent ResourceManager assertion fails. | ||
| _non_cpu_types = {g.device_type.upper() for g in device_groups} - {'CPU'} | ||
| _resources: dict[str, float] = {} | ||
| if len(_non_cpu_types) == 1: | ||
| _dev_type = next(iter(_non_cpu_types)) | ||
| if _dev_type == 'NPU': | ||
| try: | ||
| import torch | ||
| _npu_count = torch.npu.device_count() if hasattr(torch, 'npu') else 0 | ||
| if _npu_count > 0: | ||
| _resources['NPU'] = float(_npu_count) | ||
| except Exception: | ||
| pass |
There was a problem hiding this comment.
NPU 资源探测逻辑抽取为 RayHelper._get_ray_custom_resources，使 initialize() 只负责 Ray 初始化。
该方法仅在 device group 为 NPU 时通过 torch.npu.device_count() 注册 {"NPU": count}；GPU 继续由 Ray 原生自动发现，因此不会引入额外的 GPU custom resource。对于 CPU、GPU 或当前不支持的混合设备组，方法返回空字典，保持原有行为。
这里的资源注册用于当前进程启动本地 Ray 的场景；若连接已启动的 Ray 集群，则仍需在各节点启动 Ray 时注册 NPU resource。
| gkd_temperature: float = 1.0 | ||
| gkd_topk: int = 64 | ||
| router_replay_mode: Literal['disabled', 'R2', 'R3'] = 'disabled' | ||
| ppo_epochs: int = 4 |
There was a problem hiding this comment.
是不是可以复用 TrainArgs.num_train_epochs
There was a problem hiding this comment.
已经修改成使用TrainArgs.num_train_epochs
| @@ -0,0 +1 @@ | |||
| * text=auto | |||
There was a problem hiding this comment.
这个改动纯粹是行尾规范化，不影响任何运行逻辑，整个文件只有一行：
作用：
- 匹配所有文件；
text=auto 让 Git 自动区分文本/二进制文件。文本文件入库时统一按 LF 存储（无论提交者是什么平台），检出时再按当前平台转回（Windows 用 CRLF，Linux/macOS 用 LF）；二进制文件完全不碰。
为什么加：
我在 Windows/WSL 下开发，而 CI 跑在 Linux 上。不做规范化的话，只改一行代码就可能因为行尾（CRLF vs LF）变化导致整个文件显示成全量修改，污染 diff、干扰 review。
同时避免把 CRLF 意外提交进仓库，防止 Linux CI 上出现由 \r 引起的偶发失败（比如 shell/python 文件）。
它只影响 Git 对后续提交的行尾处理，不改变任何代码、依赖或行为。如果维护者认为这个 PR 不需要包含该规范化，我可以移除。
…ray NPU resource detection
| return LossOutput(loss=loss, num_tokens=0) | ||
|
|
||
|
|
||
| class PPOLoss(GRPOLoss): |
There was a problem hiding this comment.
不需要重写_aggregate_loss 吗
There was a problem hiding this comment.
确认后，PPO 的 policy loss 更适合默认按所有有效 response token 直接求平均，即 token-mean。这样每个 action token 权重一致，也与 PPOValueLoss 当前的 token-level 聚合方式保持一致。
同时在 PPOLoss 中保留了 seq-mean-token-mean 模式，即先对每条 response 内的 token loss 求平均，再对 batch 求平均。该模式能保证不同长度的 response 具有相同总权重，但会使短 response 中单个 token 的权重更大。
目前这个选择仅用于 PPO，因此没有添加到公共 CLI 参数中，而是在 PPO cookbook 中显式指定 loss_agg_mode='token-mean'。等后续其他算法也需要统一选择聚合方式时，再考虑将其抽象为公共配置。两种聚合模式均已补充单元测试。
3 participants
PR type
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:forward_only(disable_lora=True)(no extra ref GPUs)ppo_epochsmini-batch updates over fixed rolloutssrc/twinkle/advantage/gae.py—GAEAdvantage, token-level GAE with:logp_old - logp_ref) withkl_coefsrc/twinkle/loss/value.py—PPOValueLoss, clipped value objectivemax((v-r)², (clip(v,v_old)-r)²)src/twinkle/model/transformers/value_model.py—TransformersValueModel, causal-LM backbone with ahidden_size -> 1value head replacing the LM headsrc/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 statisticsFramework generalization (backward-compatible)
ModelOutputgains avaluesfield; forward/forward_only can emitvalueswhen the loss declaresrequire_valuesprocessorpacking,nccl_safewrapper, andloss/basehandle the newvaluestensor--critic-model-gpus,--ppo-epochs,--gamma,--gae-lambda,--kl-coef,--normalize-advantages,--critic-learning-rate,--value-clipTests
tests/advantage/test_gae.py— single/multi-step GAE, terminal bootstrap, padding mask, normalization, sparse reward + KL shapingtests/loss/test_ppo.py—PPOLossparity withGRPOLoss, clipped/unclipped value loss, ragged token alignmenttests/model/test_value_model.py— value-head forward/backward, save/loadtests/cli/test_cli.py— new PPO CLI argsDocs
Results
Usage