ms-swift GRPO 多任务训练指南用任务列 分任务奖励函数实现 Math/Code 混合强化学习【免费下载链接】swiftUse PEFT or Full-parameter to CPT/SFT/DPO/GRPO 600 LLMs (Qwen3.6, DeepSeek-V4, GLM-5.1, InternLM3, Llama4, ...) and 300 MLLMs (Qwen3-VL, Qwen3-Omni, InternVL3.5, Ovis2.5, GLM4.5v, Gemma4, Llava, Phi4, ...) (AAAI 2025).项目地址: https://gitcode.com/GitHub_Trending/swift1/swift本篇技术指南基于 SWIFTms-swift仓库中的 GRPO 开发者文档 Multi-Task Training讲解如何在一次 GRPO 训练中混合不同领域的数据如数学与代码任务通过在数据集里添加task任务标记列编写只对目标任务计分、对非目标任务返回 None的奖励函数从而让不同任务各自使用最合适的奖励逻辑。读完本文你将掌握多任务 GRPO 数据集的构造方式、按任务分发奖励的插件写法以及该机制在swift/rl_core/grpo_algorithm.py与swift/rlhf_trainers/grpo_trainer.py中的底层实现链路。为什么需要多任务训练GRPO 的奖励信号来自奖励函数或奖励模型。当训练集只包含单一任务例如只有数学题时一个accuracy奖励函数足以胜任但当同一个数据集里同时包含数学题和编程题时正确答案的判定标准完全不同数学答案通常用math_verify之类的符号化校验SWIFT 内置的accuracy即基于此而代码任务更依赖单元测试通过率、沙箱执行等判定方式。用同一个奖励函数去评所有样本要么数学侧不准要么代码侧失效。SWIFT 官方 FAQ 对这一问题给出的答案就是本文档本身若训练集包含不同任务请参考 Multi-Task Training 文档。核心方案只有两句话在数据集中加一列如task标记任务类型在奖励函数或奖励模型插件中读取这一列对属于当前任务的样本计算奖励对不属于的样本返回None。数据集构造添加 task 任务标记列按 multi_task.md 的示例一个同时包含数学和编程任务的 GRPO 数据集可以写成如下 JSON[ {query: Solve the equation x 2 5, solution: 3, task: math}, {query: Write a function to calculate the Fibonacci sequence, solution: xxx, task: code}, {query: What is the integral of x^2?, solution: xxx, task: math}, {query: Implement a sorting algorithm in Python, solution: xxx, task: code} ]关键点query/solution之外的额外列不会进入模型 prompt而是作为额外列原样保留供奖励函数读取。关于各列的具体语义query、response、solution等如何映射到messages可参见 自定义数据集文档。task列的取值math、code等由你自己约定奖励函数里只需按相同的字符串判断即可。注意文档中的提醒与模型输入相关的列如query、response会被转换成messages键数据集中的原始 assistant 回复会被丢弃因此像solution这类需要保留的信息应放在额外列里。编写按任务分发的奖励函数multi_task.md 给出的核心思路是数据集的列会被完整传入奖励函数因此可以在函数签名里直接声明task参数对每个样本判断它属于哪个任务属于则计算真实奖励不属于则返回None。原文给出的示例如下原文风格保留from swift.rewards import ORM, orms import random # Math-specific reward function class MathRandomReward(ORM): def __call__(self, completions, task, **kwargs): rewards [] for completion, t in zip(completions, task): if t math: # Implement math accuracy logic reward random.random() rewards.append(reward) else: # Return None for non-math tasks rewards.append(None) return rewards # Coding-specific reward function class CodeRandomReward(ORM): def __call__(self, completions, task, **kwargs): rewards [] for prompt, completion, t in zip(prompts, completions, task): if t code: # Implement coding accuracy logic reward random.random() rewards.append(reward) else: # Return None for non-coding tasks rewards.append(None) return rewards orms[math_reward] MathRandomReward orms[code_reward] CodeRandomReward说明以上代码忠实保留原文档示例。实际落地时有两处建议修正/替换CodeRandomReward.__call__中引用了未声明的prompts实现时建议只保留completions与task两个参数如数学示例所示示例中的random.random()是占位逻辑真实训练应替换为具体判定数学侧可直接复用 SWIFT 内置的accuracy奖励基于math_verify见 orm.py 中的MathAccuracy代码侧可参考仓库自带的沙箱执行示例 plugin.py 中的CodeRewardE2B 沙箱或CodeRewardByJudge0Judge0 端点。最后将奖励函数注册进orms字典如orms[math_reward] MathRandomReward注册名即可在训练参数--reward_funcs中引用。None 的语义非本任务样本不参与该奖励列对不属于当前任务的样本返回None是整套机制的枢纽。从源码看这一行为在 compute_rewards_per_func 中有明确处理每个奖励函数的输出会被逐元素转换——output reward_func(completions, **reward_kwargs) output [reward if reward is not None else torch.nan for reward in output] rewards_per_func[:, i] torch.tensor(output, dtypetorch.float32, devicedevice)也就是说None会被写入NaN形成一张[样本数, 奖励函数数]的稀疏奖励矩阵。后续在计算 advantage 时SWIFT 按列加权求和并跳过NaN因此数学奖励列只对task math的样本生效代码奖励列只对task code的样本生效两者互不干扰。同一函数里还有一道防线如果某个样本被所有奖励函数都返回None整行全NaN训练器会打印一条告警提示至少要让一个奖励函数对该样本给出有效奖励——这正是多任务分发的正确性约束每个样本必须恰好被某个奖励函数认领。数据集列是如何流进奖励函数的列会被传给奖励函数并非约定俗成而是有明确实现链路。在 GRPOSample.to_reward_row 中每个样本被压平成奖励函数可见的字典其中extra字段即数据集中不属于协议字段的额外列如task、solution会被展平到顶层随后 compute_rewards_per_func 用RowPreprocessor.rows_to_batched把逐行字典转成批量的键值对task变成与completions等长的列表与completions、trainer_state一起作为 kwargs 传给每个奖励函数。这也解释了为什么可以在__call__签名里直接写def __call__(self, completions, task, **kwargs)——参数名与列名一致即可自动绑定。补充两点实用信息若不想显式声明列名也可以从kwargs中取值task kwargs.get(task)训练步数则可通过kwargs.get(trainer_state).global_step获取见 Reward Function 文档。与模型输入相关的列会被转成messages因此奖励函数里拿到的问题文本应通过messages列访问query这类原始列名可能不再存在。训练配置external_plugins reward_funcs自定义奖励函数放在外部插件文件中通过--external_plugins加载、--reward_funcs指定参考仓库示例脚本 run_external_reward_func.shswift rlhf \ --rlhf_type grpo \ --model Qwen/Qwen2.5-3B-Instruct \ --external_plugin examples/train/grpo/plugin/plugin.py \ --reward_funcs math_reward code_reward \ --dataset 你的多任务数据集 \ --num_generations 8 \ --learning_rate 1e-6 \ ...相关参数在 GRPOConfig 中的定义参数说明默认值reward_funcs奖励函数名列表内置可选accuracy、format、cosine、repetition、soft_overlong见 orm.py自定义函数需在orms中注册[]reward_weights各奖励源的权重长度须等于奖励函数数外加外部奖励模型None时全部等权1.0None多任务场景下reward_funcs通常就是按任务拆分的多个函数如math_reward code_reward。需要留意两点组内同质性GRPO 的 advantage 是在同一 prompt 的num_generations条生成内做组内标准化奖励只作用在同任务样本上不会破坏该性质但应避免同一 prompt 的不同 completion 因任务判定不一致而部分得None。每步都要雨露均沾若某个 batch 恰好全是math数据则code_reward列整列为NaN且该列贡献为 0这是预期行为但若整行全NaN没有任何奖励函数认领该样本说明任务判定逻辑有漏洞会触发前述告警。结合内置奖励与奖励模型的多任务组合多任务机制不限于两个自定义 ORM内置奖励可直接参与分发例如数学任务直接复用内置accuracy要求安装math_verify实现见 MathAccuracy无需自己写判定逻辑代码任务再配一个沙箱执行类奖励。异步奖励兼容若某个任务的奖励涉及 I/O沙箱、API 调用可继承AsyncORM训练器会自动用asyncio.gather并行执行见 reward_function.md 与 compute_rewards_per_func 中async_indices分支仓库中的CodeReward示例本身就用了 E2B 异步沙箱。奖励模型插件同样适用文档提到reward functionor reward model plugin——在 GRPOTrainer._prepare_rewards 中--reward_model会通过rm_plugins注册为额外的奖励源。你同样可以让自定义 RM 插件在__call__里读取task列对非目标任务样本返回None奖励模型的输出在 compute_rewards_per_func 中同样经过None - NaN转换。小结SWIFT 的 GRPO 多任务训练方案可以归纳为三步数据侧为数据集添加task标记列math/code等自定义取值奖励侧按任务各写一个 ORM或 RM 插件签名里声明task参数对本任务样本计算真实奖励对非本任务样本返回None配置侧用--external_plugins加载插件文件--reward_funcs列出所有任务奖励函数必要时用--reward_weights调权。底层由 GRPOSample.to_reward_row 将数据集列透传为奖励函数的 kwargs由 compute_rewards_per_func 完成None - NaN的稀疏奖励矩阵组装最终保证每个样本只被对应的任务奖励打分。这套机制使一次 GRPO 训练即可覆盖多个领域的可验证奖励无需为每个任务单独起一轮训练。【免费下载链接】swiftUse PEFT or Full-parameter to CPT/SFT/DPO/GRPO 600 LLMs (Qwen3.6, DeepSeek-V4, GLM-5.1, InternLM3, Llama4, ...) and 300 MLLMs (Qwen3-VL, Qwen3-Omni, InternVL3.5, Ovis2.5, GLM4.5v, Gemma4, Llava, Phi4, ...) (AAAI 2025).项目地址: https://gitcode.com/GitHub_Trending/swift1/swift创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考