ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

用 GMPO 训练更稳定的 GRPO:trl 中几何均值策略优化的原理、配置与源码解析

用 GMPO 训练更稳定的 GRPO:trl 中几何均值策略优化的原理、配置与源码解析 用 GMPO 训练更稳定的 GRPOtrl 中几何均值策略优化的原理、配置与源码解析【免费下载链接】trlTrain transformer language models with reinforcement learning.项目地址: https://gitcode.com/GitHub_Trending/tr/trlGMPOGeometric-Mean Policy Optimization是 trl 在trl.experimental.gmpo模块中提供的一种 GRPO 变体它用 token 级重要性比率的几何均值取代 GRPO 的算术均值从而对离群比率不敏感、策略更新更稳定并允许使用远比 GRPO/DAPO 更宽的裁剪范围。本文基于 docs/source/gmpo.md 展开结合 GMPOTrainer 实现 与 GMPOConfig 配置 及对应的 单元测试讲清楚 GMPO 的目标函数、两个核心超参数epsilon/epsilon_high在 log 空间中的真实含义、完整的训练代码示例以及源码中裁剪与几何均值归一化的具体实现。读完你可以直接在 trl 中配置并跑通一次 GMPO 训练也能理解其与 GRPO 的精确差异。GMPO 的核心思想从算术均值到几何均值在论文Geometric-Mean Policy Optimization论文编号 2507.20673中作者提出的核心改动是GRPO 在聚合 token 级重要性比率时计算的是算术均值而 GMPO 计算的是几何均值。为什么这个改动有意义算术均值对极端值离群比率非常敏感——一个 token 出现极大的 importance ratio 就会把整条序列的目标拉偏迫使策略更新过度激进。而几何均值对离群值天然不敏感它本质上是对数空间中的平均后再指数化个别极端 token 的影响被对数运算平滑掉。因此策略更新更加稳定可以容忍更宽的裁剪范围而不至于训练崩溃裁剪按 token 进行、在log 空间执行且按 advantage 的符号做单侧裁剪即 PPO 标准 trust region 的min技巧关键的是裁剪发生在几何均值聚合之前。在 trl 中这一思想落地为trl.experimental.gmpo下的两个公开类GMPOTrainer与GMPOConfig见 gmpo/init.py。快速上手一次完整的 GMPO 训练官方文档给出的用法非常简洁核心代码如下from trl.experimental.gmpo import GMPOConfig, GMPOTrainer training_args GMPOConfig( epsilon0.4, # log 空间裁剪范围 - 比率被裁剪到 (exp(-0.4), exp(0.4))论文 Sec. 4 推荐值 beta0.0, ) trainer GMPOTrainer( modelQwen/Qwen3-0.6B, reward_funcs..., train_dataset..., argstraining_args, ) trainer.train()其中GMPOConfig与GRPOConfig的构造签名基本一致model、reward_funcs、train_dataset、args等均沿用 GRPO 的约定reward_funcs可以传入内置奖励函数或自定义奖励函数trainer.train()触发完整训练循环采样、奖励计算、优势估计、策略更新。对照仓库中的 测试用例 tests/experimental/test_gmpo_trainer.py一个可以直接运行的配置示例是training_args GMPOConfig( output_dirtmp_dir, learning_rate0.1, # 测试注释梯度偏小默认学习率可能停滞故调高 per_device_train_batch_size3, # 减小 batch 以降低显存占用 num_generations3, # 每个 prompt 采样的生成数默认 8 max_completion_length8, # 生成的最大长度 num_iterations2, # 每批迭代次数大于 1 时重要性采样权重非零 report_tonone, ) trainer GMPOTrainer( modeltrl-internal-testing/tiny-Qwen2ForCausalLM-2.5, reward_funcstrl-internal-testing/tiny-Qwen2ForSequenceClassification-2.5, argstraining_args, train_datasetdataset, ) trainer.train()测试中还覆盖了多种eval_dataset传入形态单数据集、DatasetDict、字典以及多轮对话数据集conversational_prompt_only的训练说明 GMPO 的数据流与 GRPO 完全兼容。GMPOConfig继承 GRPOConfig只改裁剪语义GMPOConfig定义在 trl/experimental/gmpo/gmpo_config.py它的类文档写得很明确继承GRPOConfig的全部参数只改变裁剪范围的语义与默认值。epsilon 与 epsilon_highlog 空间中的裁剪边界这是 GMPO 与 GRPO 最本质的差异所在参数GMPO 默认值GRPO 默认值含义epsilon0.40.2下界裁剪值以 log 空间表达token 重要性比率的下界为exp(-epsilon)epsilon_highNone回退为epsilon的值None上界裁剪值以 log 空间表达上界为exp(epsilon_high)也就是说在 GRPO 中epsilon0.2直接作为比率空间线性空间的裁剪边界而在 GMPO 中由于裁剪发生在 log 空间实际的比率裁剪区间是(exp(-epsilon), exp(epsilon_high))当epsilon_high保持None时上下界对称即(exp(-0.4), exp(0.4))。测试test_default_epsilon_is_log_space正是验证这一点def test_default_epsilon_is_log_space(self): # GMPO expresses the clip range in log space; default is the papers (exp(-0.4), exp(0.4)). args GMPOConfig(dummy) assert args.epsilon 0.4 # epsilon_high is inherited from GRPOConfig and defaults to None, so the range is symmetric. assert args.epsilon_high is None注意exp(0.4) ≈ 1.49而上界exp(0.2) ≈ 1.22、exp(0.28) ≈ 1.32DAPO 推荐值可见 GMPO 允许的比率上界明显更宽——论文推荐这个更宽的范围正是为了鼓励探索而几何均值的稳定性保证了宽范围不会导致训练发散。继承自 GRPOConfig 的其他重要参数由于GMPOConfig全量继承GRPOConfig见 trl/trainer/grpo_config.py以下参数同样生效betaKL 系数默认0.0。为0时不加载参考模型省显存、提速DeepSeek-R1 论文使用0.001。GMPO 中若beta ! 0KL 项按序列平均见下文源码解析num_iterations每个 batch 的迭代次数算法中的 μ默认1num_generations每个 prompt 的采样数默认8要求有效 batch size 能被其整除max_completion_length生成最大长度默认512temperature/top_p/top_k/min_p采样参数use_vllm/vllm_modeserver或colocate是否用 vLLM 加速生成reward_weights/multi_objective_aggregation多奖励聚合配置remove_unused_columns默认False自定义奖励函数需要额外列时保持False。beta0.0与epsilon0.4即官方文档示例中的两个显式参数——前者关闭 KL 正则、后者设定 log 空间裁剪下界。源码级解析GMPOTrainer._compute_loss 的五个步骤GMPOTrainer定义在 trl/experimental/gmpo/gmpo_trainer.py其类文档明确说明相对GRPOTrainer的唯一改动是_compute_loss其余部分生成、奖励计算、权重同步、指标日志原样继承。因此理解 GMPO 的关键就是逐行读懂_compute_loss见 gmpo_trainer.py#L43-L155。第 1 步计算 per-token log 概率与熵拼接 prompt 与 completion仅保留 completion 段的 logitslogits_to_keep调用_get_per_token_logps_and_entropies得到每个位置的对数概率per_token_logps与熵entropies。注意该方法支持多模态输入pixel_values、image_grid_thw、image_sizes等参数均有透传说明 GMPO 沿用了 GRPO 的视觉语言模型支持。若top_entropy_quantile 1.0还会构造高熵 token 掩码entropy_mask把低熵 token 从目标中剔除只让高熵 token 参与梯度。当num_iterations 1且steps_per_generation gradient_accumulation_steps时新旧策略的 logps 相同实现上直接用per_token_logps.detach()跳过前向计算以省算力。第 2 步log 空间中的逐 token 双向裁剪log_ratio per_token_logps - old_per_token_logps # per-token log 重要性比率 clamped_log_ratio torch.clamp(log_ratio, min-self.epsilon_low, maxself.epsilon_high)裁剪直接作用于对数比率log 空间对应比率空间的(exp(-epsilon), exp(epsilon_high))区间数值上比在比率空间裁剪更稳定。第 3 步按 advantage 符号做单侧裁剪PPO trust regionclipped_log_ratio torch.where( advantages_col 0, torch.minimum(log_ratio, clamped_log_ratio), torch.maximum(log_ratio, clamped_log_ratio), )这是 PPO 经典 min 技巧的 log 空间写法advantage 为正的 token 只限制上界防止过度上调该 token 的概率advantage 为负的 token 只限制下界防止过度下调。单侧裁剪在几何均值之前逐 token 完成——这正是文档强调的 one-sided per the advantage sign, before the geometric mean is taken。第 4 步几何均值聚合与逐序列损失log_importance_weights (clipped_log_ratio * seq_mask).sum(-1) / seq_mask.sum(-1).clamp(min1.0) coef torch.exp(log_importance_weights) # 序列级几何均值重要性权重 per_sequence_loss -coef * advantages先对有效 token 的裁剪后 log 比率求平均对数空间平均再指数化得到几何均值coef代码注释特别指出1/|o_i|指数正是几何均值的归一化因子论文的消融实验表明这一归一化是必需的逐序列损失为-coef * advantages对序列做普通平均得到 batch 损失token 归一化已蕴含在几何均值内部并按梯度累积步数归一化。第 5 步可选的 KL 正则与裁剪指标当beta ! 0.0时向参考模型方向施加 KL 正则且与 GMPO 的序列级目标对齐KL 同样按序列平均测试test_train_with_kl专门覆盖beta0.1路径per_token_kl torch.exp(ref_per_token_logps - per_token_logps) - (ref_per_token_logps - per_token_logps) - 1 seq_kl (per_token_kl * mask).sum(-1) / mask.sum(-1).clamp(min1.0) per_sequence_loss per_sequence_loss self.beta * seq_kl此外训练日志中会记录以下裁剪相关指标用于监控有多少 token 被推进裁剪区klbeta ! 0时、entropyclip_ratio/low_meanlog_ratio -epsilon_low且advantage 0的 token 占比clip_ratio/high_meanlog_ratio epsilon_high且advantage 0的 token 占比clip_ratio/region_mean两侧裁剪区合并占比clip_ratio/low_min/clip_ratio/high_max跨进程 gather 后的序列级极值。GMPOTrainer的_tag_names [trl, gmpo]训练完成后会自动打上这两类模型标签train、save_model、push_to_hub等方法均从GRPOTrainer继承用法与 GRPO 完全一致。训练实践建议与注意事项结合 测试用例 与源码给出几点实操要点学习率GMPO 的梯度通常偏小测试中特意将learning_rate提到0.1以避免更新停滞实际训练时如发现 loss 不下降可优先排查学习率量级。num_iterations 1时的重要性采样测试注释指出num_iterations2时重要性采样权重非零说明多迭代下old_per_token_logps才真正来自旧策略这是 GRPO 家族通用的机制。默认beta0.0不加载参考模型内存占用更小、速度更快需要约束与参考模型的偏离时再开启DeepSeek-R1 经验值0.001可作起点。epsilon/epsilon_high的非对称用法epsilon_high独立于epsilon设置时即可实现非对称裁剪如epsilon0.4, epsilon_high0.28满足对上下界不同容忍度的场景。限制测试注释明确Streaming datasets are not yet supported in GMPO流式数据集暂不支持eval_dataset支持单数据集、DatasetDict、字典三种形态。作为实验性功能GMPOTrainer位于trl.experimental.gmpo命名空间下API 属于实验性模块升级 trl 时需留意接口变更MIGRATION.md 中会有相关迁移说明。小结GMPO 与 GRPO 的差异可以浓缩为一句话GRPO 在比率空间用算术均值GMPO 在 log 空间裁剪后取几何均值。在 trl 中这只意味着换一个 Trainer 类、改一个epsilon默认值——生成、奖励、权重同步全部复用 GRPO 的成熟实现而目标函数的稳定性收益则由 gmpo_trainer.py 中几十行_compute_loss精妙地承载。如果你想在 GRPO 基础上获得更稳的策略更新并让探索更充分GMPO 是一个低接入成本、高理论清晰度的选择。【免费下载链接】trlTrain transformer language models with reinforcement learning.项目地址: https://gitcode.com/GitHub_Trending/tr/trl创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表