近端策略优化 (PPO) - 亚马逊 SageMaker AI

本文属于机器翻译版本。若本译文内容与英语原文存在差异,则一律以英文原文为准。

近端策略优化 (PPO)

近端策略优化 (PPO) 是使用多个机器学习模型对模型进行训练并评分的过程。以下模型是 PPO 流程的一部分:

  • 行为者训练或策略模型:一种监督式微调 (SFT) 模型,每个时代都会进行微调和更新。更新是通过对提示进行采样、生成完成结果以及使用剪辑代理目标更新权重完成的。这可以限制单 Token 对数收益的变化,以使各策略步骤接近于上一个步骤,从而维持训练的稳定性。

  • 角色生成模型:一种生成提示完成结果或响应的模型,通过奖励模型和评估模型进行判断。该模型的权重在每个 epoch 从角色训练或策略模型中进行更新。

  • 奖励模型:一种具有冻结权重的模型,用于对演员生成模型进行评分。

  • 评论家模型:一种具有未冻结权重的模型,用于对演员生成模型进行评分。该分数通常被视为行为者在生成剩余代币时获得的总奖励的估计值。

  • 锚模型:具有冻结权重的 SFT 模型,用于计算角色列车模型和基础模型之间的 KL 差异。锚点模型可确保角色模型的更新与基础模型相比不会太剧烈。剧烈变化可能导致不稳定或性能下降。

训练数据必须采用 JSONL 格式,其中每行都包含一个代表训练示例的 JSON 对象。示例如下:

{ "turns": ["string", "string", ...], // Required "turns_to_mask": [integer, integer, ...], // Required "reward_category": "string", // Required "meta_data": {} // Optional }
  • turns是代表对话序列的对话字符串数组的数组。此行包含系统提示、用户消息和机器人响应。用户消息通常以 “Bot:” 结尾,以指示模型输出从哪里开始。例如 [["System prompt"], ["User: Question Bot:"], ["Bot response"]]

  • turns_to_mask是一个基于 0 的索引数组,用于标识哪些回合不应接收渐变更新。被屏蔽的转弯通常是系统提示和用户回合。例如,[0, 1, 3]屏蔽系统提示和用户消息(第一和第三条消息)。

  • reward_category是一个字符串,用于标识要评估模型性能的哪些方面。它用于在训练期间选择适当的奖励模型类别。奖励类别适用于以下奖励类别:defaultmathcodingifrag、和rai

  • meta_data是一个可选对象,其中包含其他上下文或真实信息。这可能包括标识符、来源信息或对话上下文。根据您的数据集需求,该结构非常灵活。

以下是示例记录:

{ "turns": ["You are a helpful AI assistant.", "User: What is ML? Bot:", "Machine learning is...", "User: Examples? Bot:", "Email spam filtering is..." ], "turns_to_mask": [0, 1, 3], "reward_category": "default", "meta_data": { "messages": [{ "role": "system", "content": "You are a helpful AI assistant." }, { "role": "user", "content": "What is ML?" }, { "role": "assistant", "content": "Machine learning is..." }, { "role": "user", "content": "Examples?" }, { "role": "assistant", "content": "Email spam filtering is..." } ] } }

奖励建模框架实现了跨不同类别目标的多维优化,以促进稳健的模型收敛性。应根据模型必须优化的任务来选择奖励类别。

我们建议您遵循以下指导方针,为您的任务选择合适的框架:

  • default:用于标准对话任务和基本交互的通用优化器。用于一般对话和讨论、基本写作任务、简单问答和非专业知识查询。

    示例如下:

    { "turns": ["Write a summary of climate change"], "turns_to_mask": [0], "reward_category": "default" }
  • math:用于数学计算和数值推理任务的专业优化器。用于解决数学问题、算术计算、代数方程、几何问题和统计分析。

    示例如下:

    { "turns": ["Calculate the derivative of x²"], "turns_to_mask": [0], "reward_category": "math" }
  • coding:专门用于编程和软件开发相关查询的类别。用于代码实现、调试帮助、算法设计、技术文档和系统架构问题。

    示例如下:

    { "turns": ["Write a function to check if a string is palindrome"], "turns_to_mask": [0], "reward_category": "coding" }
  • if: 需要精确程序执行和 step-by-step指导的任务类别。用于多步骤过程、顺序指令、复杂任务分解和过程文档。

    示例如下:

    { "turns": ["Provide steps to deploy a web application"], "turns_to_mask": [0], "reward_category": "if" }
  • rag:奖励类别,用于专门根据检索到的上下文信息回答查询的任务。当应直接从提供的参考材料中得出答复时使用,在不超出检索信息的范围的情况下综合事实内容,确保答案基于所提供的背景而不是常识。

    示例如下:

    { "turns": ["The Synthesis Report integrates findings from all six IPCC assessment cycles, revealing that global surface temperature has increased 1.1°C from 1850-1900 to 2011-2020, with human activities unequivocally identified as the cause of this warming. Alarmingly, current policies put the world on track for 3.2°C warming by 2100. The document identifies 5 key climate system "tipping points" approaching and emphasizes that greenhouse gas emissions must decline 43% by 2030 (compared to 2019 levels) to limit warming to 1.5°C. Climate-related risks will escalate with every increment of warming, with loss and damage disproportionately affecting vulnerable populations. Despite some progress, climate adaptation remains uneven with significant gaps, and financial flows continue to fall below levels needed for mitigation goals.", "What were the key findings of the latest IPCC climate report?"], "turns_to_mask": [0, 0], "reward_category": "rag" }
  • rai: 一种奖励类别,用于奖励需要应用负责任的人工智能原则,例如公平、透明和道德。用于评估人工智能系统中的潜在偏见、确保隐私考虑、解决道德困境和促进包容性设计原则。

    示例如下:

    { "turns": ["Identify potential bias concerns when developing a loan approval algorithm and suggest mitigation strategies"], "turns_to_mask": [0], "reward_category": "rai" }
掩盖转弯

在训练数据集中,该turns_to_mask参数对于控制哪些对话在训练期间接收梯度更新至关重要。这个索引数组决定了模型应该学会生成对话的哪些部分,而哪些部分应仅被视为上下文。适当的屏蔽可确保模型学习适当的响应模式,同时避免根据系统提示或用户输入进行可能降低性能的训练。

我们建议使用以下蒙版指南:

  • 始终屏蔽索引 0-系统提示不应接收渐变更新。

  • 始终屏蔽用户回合-防止模型学习生成用户输入。

  • 模式一致性-对相似的对话结构使用相同的屏蔽模式,例如 (0、1、3、5) 用于多回合对话。

  • 选择性训练-掩盖早期的机器人反应,将训练重点放在改进的最终反应上。

  • Chain-of-thought 保存-在训练推理序列时,只有掩码系统和用户才会转动。

  • 质量过滤-屏蔽低质量的助手响应以防止性能下降。

  • 情境优化-确保掩盖转弯不会移除后续响应所需的基本上下文。

有效屏蔽的关键是监控训练指标和验证性能,以确定您的屏蔽策略是否保留了必要的上下文,同时将梯度更新重点放在所需的模型输出上。

启用 KL 背离损失

要启用 KL Divergence 损失,需要启用锚服务器以计算当前策略与原始分配的差异。需要指定 KL 损失类型,并且系数必须是非零的值。较高的系数值有助于模型不会与原始策略有太大偏差,从而减少总体性能的变化。较低的系数值允许与先前策略的偏差更大,从而提高目标指标的性能,但会影响总体绩效。

ppo_anchor: max_length: 8192 trainer: num_nodes: ${recipes.run.cm_replicas} model: global_batch_size: 32 ppo_actor_train: model: ######## Use KL in actor loss ######## kl_loss_type: low_var_kl kl_loss_coeff: 0.1 ######## Use KL in reward model ###### kl_reward_penalty_coeff: 0.1
学习率

批评者和政策模型的学习率可以调整,3e-6是默认的平衡选择。较高的学习率通常会导致训练不稳定,这可以通过KL差异峰值和不稳定的政策行为来识别。较低的学习率可能会导致融合问题和学习缓慢,这表现为奖励停滞不前和政策更新很少。定期监控 KL 差异、奖励分数和价值损失有助于确定是否在训练期间调整学习率。

ppo_critic: model: optim: lr: 3e-6 ppo_actor_train: model: optim: lr: 3e-06
全局批次大小

全球批量大小会显著影响 Amazon Nova 中的 PPO 性能,较大的批次通常可以提高训练稳定性和梯度估计,同时实现更高效的并行处理。但是,非常大的批量大小会导致收益减少,并且可能会受到可用内存的限制,因此需要仔细平衡学习率和其他超参数。

ppo_actor_train: model: global_batch_size: 160

可用于 PPO 调整的 Amazon Nova 参数包括:

  • 运行配置

    • actor_train_replicas:用于演员训练模型的计算实例数量。可用值因所选型号而异。亚马逊 Nova Micro 支持 1 或 2 个副本。Amazon Nova Lite 支持 1、2 或 4 个副本。Amazon Nova Pro 支持 3、6 或 12 个副本。

    • rm_replicas:用于奖励模型的计算实例数量。我们建议您为任何模型尺寸使用一个副本。

    • cm_replicas:用于批评模型的计算实例数量。我们建议您为任何模型尺寸使用一个副本。

    • actor_generation_replicas:用于生成角色的计算实例的数量。可用值因所选型号而异。亚马逊 Nova Micro 支持 1 个副本。Amazon Nova Lite 支持 1 或 2 个副本。亚马逊 Nova Pro 支持 1 或 2 个副本。

    • am_replicas:用于锚模型的计算实例数。我们建议您为任何模型尺寸使用一个副本。

  • 演员训练配置(策略配置)

    • max_steps: 微调或训练演员训练模型的最大步骤数。在这里,一个步骤被定义为推出,然后使用样本global_batch_size数量训练演员训练模型。一个时代被定义为。global_batch_size * trajectory_buffer_scale

      此处选择的值将根据您的用例和数据集的复杂性而有所不同。我们建议从 65 个周期或 520 步开始,即周期数乘以的值。trajectory_buffer_scale但是,有些任务需要更长的 PPO 训练时间才能达到同样的性能。

      对于 PPO,训练指标,例如来自 ml-flow 控制台的饱和奖励模型分数和平均动作时长,可以帮助确定最佳评估点。

    • actor_model_max_length:发送到角色生成组件以生成完成的输入数据的最大长度。

    • reward_model_max_length:发送到奖励服务器以获得完成分数的输入数据的最大长度。

    • trajectory_buffer_scale:此缓冲区表示在更新权重和生成新部署之前,使用旧的 actor train(策略)模型生成的部署数量。支持的值为 1、2、4、8 和 16。

      如果trajectory_buffer_scale为 1,则培训是根据政策进行的。这意味着部署是使用最新的模型权重生成的,但吞吐量会受到影响。如果是 16,则模型稍微偏离政策,但吞吐量会更高。我们建议每个型号从 8 开始。

    • kl_reward_penalty_coeff:这是KL差异术语,可确保更新不会太激烈,并且政策不会从基础模型或SFT模型中起草。

    • kl_loss_coeff:此值控制 KL 差异惩罚在多大程度上影响 PPO 中的总体训练目标。

    • kl_loss_type:此值指定如何计算当前保单分配和参考保单分配之间的差异。kl_loss_types可用的有kl(标准 KL 散度)、mse(均方误差)、abs(对数概率之间的绝对差)和low_var_kl(低方差 KL 近似值)。

    • model.clip_ratio:PPO 中的角色剪辑比率 (ε) 是一个超参数,用于限制每次更新期间策略可以更改的程度。

    • model.optim.lr:演员模型中用于代理模型损失训练的学习速率。

    • model.lam:优势估算过程的一部分。λ越高,对长期奖励的权重越大,但方差越大,而 λ越低,则更多地关注方差较小但偏差越大的即时奖励。

    • model.ent_coeff: PPO中的熵损失通过在政策变得过于确定性时对其进行惩罚来鼓励探索(也就是说,始终充满信心地选择相同的行动)。

  • 奖励模型配置

    • global_batch_size:使用奖励模型对完成量进行评分的批次大小。如果大ppo_actor_train.model.global_batch_sizeppo_reward.model.global_batch_size,则分多个批次处理它们。请注意,ppo_actor_train.model.global_batch_size % ppo_reward.model.global_batch_size必须等于 0。

    • max_length:奖励模型的最大上下文长度。这应该与相同ppo_actor_train.model.max_length

  • 评论家模型配置

    • global_batch_size:批评模型值的批次大小。批评者模型将为参与者模型提供的响应中的每个代币提供价值估计。批量大小用于推理和训练。

      请注意,ppo_actor_train.model.global_batch_size % ppo_critic.model.global_batch_size必须等于 0 和ppo_actor_train.model.global_batch_size * ppo_actor_train.model.trajectory_buffer_size % ppo_critic.model.global_batch_size == 0

    • max_length:批评者模型的最大上下文长度。这应该与相同ppo_actor_train.model.max_length

    • model.optim.lr:演员模型中用于代理模型损失训练的学习速率。

  • 锚模型配置

    • global_batch_size:生成冻结的 SFT 或锚模型的日志的批次大小。请注意,ppo_actor_train.model.global_batch_size % ppo_anchor.model.global_batch_size必须等于 0。

    • max_length:奖励模型的最大上下文长度。这应该与相同ppo_actor_train.model.max_length

  • Actor 生成模型配置

    • actor_model_max_length:actor 模型生成组件的最大上下文长度。这应该与相同ppo_actor_train.model.max_length

PPO 食谱

以下是 PPO 的配方。

## Run config run: name: ndry-ppo-pro model_type: amazon.nova-pro-v1:0:300k model_name_or_path: nova-pro/prod data_s3_path: s3://testing/train.jsonl # Your training data S3 path actor_train_replicas: 6 # Actor train model replicas rm_replicas: 1 # Reward model replicas cm_replicas: 1 # Critic model replicas actor_generation_replicas: 2 # Actor generation model replicas am_replicas: 1 # Anchor model replicas ## Training config for each PPO component ppo_reward: max_length: 8192 # model architecture max length trainer: num_nodes: ${recipes.run.rm_replicas} model: global_batch_size: 16 ppo_critic: max_length: 8192 trainer: num_nodes: ${recipes.run.cm_replicas} model: global_batch_size: 16 optim: lr: 3e-6 name: distributed_fused_adam adam_w_mode: true eps: 1e-06 weight_decay: 0.0 betas: - 0.9 - 0.999 ppo_anchor: max_length: 8192 trainer: num_nodes: ${recipes.run.am_replicas} model: global_batch_size: 16 ppo_actor_generation: actor_model_max_length: 8192 trainer: num_nodes: ${recipes.run.actor_generation_replicas} ppo_actor_train: max_length: 8192 max_steps: 520 # Stopping criteria Desired epoch num * trajectory_buffer_scale actor_model_max_length: 8192 # truncate input data to max length reward_model_max_length: 8192 # truncate input data to max length trajectory_buffer_scale: 8 trainer: num_nodes: ${recipes.run.actor_train_replicas} model: global_batch_size: 160 ent_coeff: 0 clip_ratio: 0.2 lam: 1 kl_loss_coeff: 0.0 kl_loss_type: low_var_kl kl_reward_penalty_coeff: 0.0 hidden_dropout: 0.0 # Dropout probability for hidden state transformer. attention_dropout: 0.0 # Dropout probability in the attention layer. ffn_dropout: 0.0 # Dropout probability in the feed-forward layer. optim: lr: 3e-06 name: distributed_fused_adam # only this one is available for p0. adam_w_mode: true eps: 1e-08 weight_decay: 0.0 betas: - 0.9 - 0.999
限制

PPO 有以下限制:

  • 中间检查点不会保存以供评估,也无法从中间检查点恢复。仅保存最后一个检查点。

  • 不支持多模式数据集。

  • 训练作业不会自动停止。你必须使用 HyperPod CLI 停止作业。

  • MLflow 不支持日志记录。

  • 不支持评论家训练指标 TensorBoard。

  • 要调整超参数,请按照选择超参数中的指导进行操作。