本次实践将监督微调(SFT)、奖励模型(RM)以及使用PPO的RL微调的各个阶段整合到一个功能性的、简化的端到端循环中。此处的目的并非达成最先进的对齐效果,而是演示这些组成部分如何在一个训练过程中相互作用,并使用Hugging Face TRL库等工具。我们假设您具备以下条件:一个适用于SFT的预训练语言模型(例如 gpt2、distilbert-base-uncased)。为简单起见,在此最小示例中,我们甚至可能直接使用基础模型作为起始的“SFT”策略。一个与基础LM架构兼容的预训练奖励模型。该模型应接收一个(query, response)对(或仅接收响应,具体取决于其训练方式),并输出一个标量分数。一个已安装 transformers、torch(或 tensorflow)和 trl 的Python环境。配置组成部分首先,我们需要加载模型并配置PPO训练器。我们将使用占位符名称;请将其替换为您的实际模型路径或Hugging Face标识符。import torch from transformers import AutoTokenizer, AutoModelForCausalLMWithValueHead, pipeline from trl import PPOTrainer, PPOConfig, AutoModelForCausalLMWithValueHead # 1. 配置 - 用于演示的最小PPO设置 ppo_config = PPOConfig( model_name="gpt2", # 或者您的特定SFT模型路径 learning_rate=1.41e-5, batch_size=4, # 用于演示的小批量大小 mini_batch_size=2, gradient_accumulation_steps=1, log_with="tensorboard", # 可选:用于日志记录 kl_penalty="kl", # 使用KL惩罚 target_kl=0.1, # 目标KL散度 init_kl_coef=0.2, # 初始KL系数 adap_kl_ctrl=True, # 使用自适应KL控制 ppo_epochs=4, # 每批次的优化周期数 seed=0, ) # 2. 加载模型和分词器 # 策略模型(Actor/Critic):从SFT/基础模型初始化 # AutoModelForCausalLMWithValueHead 结合了LM头和价值头 policy_model = AutoModelForCausalLMWithValueHead.from_pretrained(ppo_config.model_name) # 参考模型(用于KL散度):保留初始策略的副本 ref_model = AutoModelForCausalLMWithValueHead.from_pretrained(ppo_config.model_name) tokenizer = AutoTokenizer.from_pretrained(ppo_config.model_name) # 确保为分词器设置了pad token if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token # 奖励模型(RM):单独加载,假设它是一个文本分类风格的pipeline # 替换! reward_model_name = "path/to/your/reward/model" # 替换! # 注意:这可能需要自定义pipeline或直接加载模型,具体取决于您的RM try: # 简化示例,假设是兼容的情感/奖励pipeline reward_pipe = pipeline("text-classification", model=reward_model_name, device=policy_model.device) print("Reward model loaded via pipeline.") # 定义一个函数来获取标量分数 def get_reward_score(texts): # 处理文本,如果需要,可能将其格式化为 (query, response) # 这高度依赖于您的奖励模型的输入格式 # 假设RM输出一个字典列表,如 [{'label': 'POSITIVE', 'score': 0.9}] results = reward_pipe(texts, return_all_scores=True) # 根据您的pipeline进行调整 # 提取所需的分数(例如,“POSITIVE”的分数或特定的分数索引) # 此提取逻辑高度依赖于您的RM的输出 scores = [] for result in results: # 示例:查找特定标签的分数,或假设第一个分数是奖励 # 根据您的奖励模型结构调整此逻辑 score = 0.0 # 默认分数 if isinstance(result, list): # 处理变化的pipeline输出 for label_score in result: if label_score['label'] == 'POSITIVE': # 示例标签 score = label_score['score'] break elif isinstance(result, dict): score = result.get('score', 0.0) # 简化回退 scores.append(torch.tensor(score, device=policy_model.device)) return scores except Exception as e: print(f"Warning: Could not load reward model pipeline '{reward_model_name}'. Using dummy rewards. Error: {e}") # 如果RM加载失败,则回退到虚拟奖励函数 def get_reward_score(texts): # 虚拟奖励:基于长度的分数(仅用于演示) return [torch.tensor(len(text) / 100.0, device=policy_model.device) for text in texts] # 3. 初始化PPOTrainer ppo_trainer = PPOTrainer( config=ppo_config, model=policy_model, ref_model=ref_model, tokenizer=tokenizer, # 在此手动循环示例中,可以省略数据集 # 如果不使用 AutoModelForCausalLMWithValueHead,则value_model需要单独设置 ) print("设置完成。已准备好运行简化的RLHF循环。")简化的RLHF循环现在,我们来执行RLHF循环的几个步骤。我们将手动提供查询、生成响应、获取奖励并执行PPO更新。# 定义一些示例查询 queries = [ "Explain the concept of KL divergence in simple terms:", "Write a short poem about a robot learning:", "What are the main stages of RLHF?", "Suggest a name for a friendly AI assistant:", ] # 对查询进行分词 query_tensors = [tokenizer.encode(q, return_tensors="pt").to(policy_model.device) for q in queries] # 策略模型的生成设置 generation_kwargs = { "min_length": -1, # 允许提前停止 "top_k": 0.0, "top_p": 1.0, "do_sample": True, "pad_token_id": tokenizer.pad_token_id, "max_new_tokens": 64, # 限制响应长度用于演示 } # 运行几个PPO步骤(例如2步) num_steps = 2 for step in range(num_steps): print(f"\n--- PPO Step {step + 1} ---") # 1. 策略执行:从策略模型生成响应 response_tensors = [] for query_tensor in query_tensors: # 生成响应;响应包含查询和生成的部分 response = ppo_trainer.generate(query_tensor.squeeze(0), **generation_kwargs) response_tensors.append(response.squeeze()) # 解码响应以进行奖励计算和日志记录 decoded_responses = [tokenizer.decode(r.squeeze(), skip_special_tokens=True) for r in response_tensors] # 2. 奖励计算:使用RM对生成的响应进行评分 # 如果需要,为奖励模型格式化文本(例如,组合查询+响应) # 此示例假设RM对包括提示在内的完整生成文本进行评分 reward_texts = decoded_responses rewards = get_reward_score(reward_texts) # 张量标量列表 # 3. PPO优化步骤 # 为 ppo_trainer.step 准备输入 # query_tensors 需要是 List[torch.Tensor] # response_tensors 需要是 List[torch.Tensor] # rewards 需要是 List[torch.Tensor](每个样本的标量奖励) stats = ppo_trainer.step(query_tensors, response_tensors, rewards) # 4. 日志记录 print(f"Query examples: {[q[:50] + '...' for q in queries]}") print(f"Response examples: {[r[len(q):][:80] + '...' for q, r in zip(queries, decoded_responses)]}") print(f"Mean reward: {torch.mean(torch.stack(rewards)).item():.4f}") if 'ppo/kl' in stats: print(f"KL Divergence: {stats['ppo/kl']:.4f}") if 'ppo/loss/policy' in stats: print(f"Policy Loss: {stats['ppo/loss/policy']:.4f}") if 'ppo/loss/value' in stats: print(f"Value Loss: {stats['ppo/loss/value']:.4f}") # 可选:如果使用TensorBoard等日志记录器,则记录详细统计信息 # ppo_trainer.log_stats(stats, queries, response_tensors, rewards) print("\n简化的RLHF循环已完成。") 数据流图下图描述了此简化循环单次迭代中的数据流:digraph RLHF_Loop { rankdir=LR; node [shape=box, style=rounded, fontname="Arial", fontsize=10]; edge [fontname="Arial", fontsize=9]; subgraph cluster_models { label = "模型"; style=filled; color="#e9ecef"; // gray policy_model [label="策略模型\n(Actor-Critic)", shape=cylinder, style=filled, color="#a5d8ff"]; // blue reward_model [label="奖励模型\n(固定)", shape=cylinder, style=filled, color="#96f2d7"]; // teal } subgraph cluster_ppo { label = "PPO训练器"; style=filled; color="#e9ecef"; // gray ppo_update [label="PPO更新步骤\n(损失计算与优化)", style=filled, color="#ffec99"]; // yellow } query [label="输入查询\n(批量)", shape=Mdiamond, style=filled, color="#fcc2d7"]; // pink response [label="生成响应", style=filled, color="#bac8ff"]; // indigo reward_score [label="奖励分数", shape=invtriangle, style=filled, color="#ffe066"]; // yellow stats [label="训练统计\n(损失, KL, 奖励)", shape=note, style=filled, color="#ced4da"]; // gray query -> policy_model [label=" 生成"]; policy_model -> response [label=" 采样文本"]; response -> reward_model [label=" 评分"]; reward_model -> reward_score; query -> ppo_update [label=" 查询输入"]; response -> ppo_update [label=" 响应输入"]; reward_score -> ppo_update [label=" 奖励信号"]; ppo_update -> policy_model [label=" 更新策略权重"]; ppo_update -> stats [label=" 记录"]; }RLHF中单个PPO步骤的数据流。查询促使策略模型生成响应,这些响应随后由奖励模型评分。PPO训练器使用查询、响应和奖励分数来计算损失并更新策略模型的参数。观察与说明在此简化执行中,您应注意以下几点:响应生成:策略模型根据输入查询生成文本。最初,这些响应反映SFT模型(或基础模型)的行为。奖励分配:奖励模型为这些响应分配分数。如果使用真实的RM,这些分数应与期望的特点(例如,有用性、无害性)相关联。如果使用虚拟奖励,它只是一个占位符信号。PPO统计数据:ppo_trainer.step函数返回统计数据。请关注mean_reward;理想情况下,如果策略学习生成RM偏好的响应,它应该呈上升趋势。请留意KL Divergence (ppo/kl),以确保策略不会与原始参考模型偏离过大,从而防止崩溃。策略和价值损失(ppo/loss/policy、ppo/loss/value)表示优化进度。策略变化:尽管在仅几个步骤和小批量的情况下变化不明显,但策略模型的参数正在更新。如果您使用更大的数据集运行更多步骤,您将开始看到生成响应的风格或内容发生变化,有望更好地与奖励模型中编码的偏好对齐。“本次实践将RLHF过程简化为其核心循环:生成、评分、更新。它突出了策略模型、奖励模型以及由训练器管理的PPO算法之间的交互点。虽然RLHF涉及更大的规模、复杂的数据处理、仔细的超参数调整和分布式训练,但这个动手示例为本章讨论的各组成部分之间的底层机制提供了切实的感受。在此基础上,您可以扩展实现、整合适当的数据集处理,并完善配置以进行更重要的对齐任务。”