企业官网建设流程全解析

grpo_from_scratch/ ├── train.py # 主训练脚本 ├── grpo_loss.py # GRPO 损失函数实现 ├── reward_fn.py # 奖励函数实现 ├── data_utils.py # 数据加载与处理 └── inference.py # 推理与评估

import torch import torch.nn.functional as F from torch import nn from typing import List, Optional import math class GRPOTrainer: """ GRPO (Group Relative Policy Optimization) 训练器 GRPO 的核心创新： 1. 使用组内标准化替代价值网络 2. 显式 KL 惩罚防止策略崩塌 3. 裁剪确保更新稳定性 """ def __init__( self, policy_model: nn.Module, # 正在训练的策略模型 ref_model: nn.Module, # 冻结的参考模型 tokenizer, beta: float = 0.04, # KL 惩罚系数 epsilon: float = 0.2, # PPO clip 范围 group_size: int = 8, # 每组采样数 learning_rate: float = 1e-6, max_grad_norm: float = 1.0, ): self.policy_model = policy_model self.ref_model = ref_model self.ref_model.eval() # 冻结参考模型 for param in self.ref_model.parameters(): param.requires_grad = False self.tokenizer = tokenizer self.beta = beta self.epsilon = epsilon self.group_size = group_size self.max_grad_norm = max_grad_norm self.optimizer = torch.optim.AdamW( self.policy_model.parameters(), lr=learning_rate, weight_decay=0.01 ) def compute_log_probs( self, model: nn.Module, input_ids: torch.Tensor, attention_mask: torch.Tensor, ) -> torch.Tensor: """ 计算 token 级别的 log 概率 返回: [batch_size, seq_len] - 每个位置的 log prob """ outputs = model(input_ids=input_ids, attention_mask=attention_mask) logits = outputs.logits # [batch_size, seq_len, vocab_size] # 计算每个 token 的 log probability log_probs = F.log_softmax(logits, dim=-1) # 收集实际 token 的 log prob token_log_probs = log_probs.gather( dim=-1, index=input_ids.unsqueeze(-1) ).squeeze(-1) # [batch_size, seq_len] return token_log_probs def compute_sequence_log_probs( self, log_probs: torch.Tensor, attention_mask: torch.Tensor, ) -> torch.Tensor: """ 计算整个序列的 log 概率（已生成部分） 返回: [batch_size] - 每个序列的总 log prob """ # 只计算生成的 token（排除 prompt 部分） # attention_mask 中 1=有效token return (log_probs * attention_mask).sum(dim=-1) def compute_kl_penalty( self, policy_log_probs: torch.Tensor, ref_log_probs: torch.Tensor, attention_mask: torch.Tensor, ) -> torch.Tensor: """ 计算 KL 散度惩罚 使用近似公式: KL ≈ exp(ref_log_prob - policy_log_prob) - (ref_log_prob - policy_log_prob) - 1 这个近似比直接的 KL 计算更稳定，且梯度特性更好 """ log_ratio = ref_log_probs - policy_log_probs # 因为我们在计算策略的 KL 惩罚 kl_div = torch.exp(log_ratio) - log_ratio - 1 # 只计算生成部分的 KL kl_div = kl_div * attention_mask return kl_div.sum(dim=-1) # [batch_size] def grpo_loss( self, input_ids: torch.Tensor, # [batch_size * group_size, seq_len] attention_mask: torch.Tensor, # 同 input_ids rewards: torch.Tensor, # [batch_size * group_size] ) -> torch.Tensor: """ GRPO 损失计算 核心流程： 1. 计算策略模型和参考模型的 log probs 2. 计算重要性采样比率 3. 组内标准化奖励 → 优势 4. 计算裁剪后的策略梯度损失 5. 加上 KL 惩罚 """ batch_size = rewards.shape[0] assert batch_size % self.group_size == 0 num_groups = batch_size // self.group_size # ---- Step 1: 计算 log probs ---- with torch.no_grad(): # 参考模型的 log probs（冻结，不计算梯度） ref_token_log_probs = self.compute_log_probs( self.ref_model, input_ids, attention_mask ) # 策略模型的 log probs（需要梯度） policy_token_log_probs = self.compute_log_probs( self.policy_model, input_ids, attention_mask ) # ---- Step 2: 序列级别的 log probs ---- # 计算生成部分的 log probs policy_seq_log_probs = self.compute_sequence_log_probs( policy_token_log_probs, attention_mask ) ref_seq_log_probs = self.compute_sequence_log_probs( ref_token_log_probs, attention_mask ) # 重要性采样比率 log_ratio = policy_seq_log_probs - ref_seq_log_probs.detach() ratio = torch.exp(log_ratio) # ---- Step 3: 组内标准化奖励 → 优势 ---- advantages = self._compute_group_advantages(rewards, num_groups) # ---- Step 4: 裁剪后的策略梯度损失 ---- pg_loss1 = -advantages * ratio pg_loss2 = -advantages * torch.clamp( ratio, 1.0 - self.epsilon, 1.0 + self.epsilon ) pg_loss = torch.max(pg_loss1, pg_loss2).mean() # ---- Step 5: KL 惩罚 ---- kl_penalty = self.compute_kl_penalty( policy_token_log_probs, ref_token_log_probs, attention_mask ).mean() # 总损失 loss = pg_loss + self.beta * kl_penalty return loss, { "pg_loss": pg_loss.item(), "kl_penalty": kl_penalty.item(), "mean_reward": rewards.mean().item(), "mean_advantage": advantages.mean().item(), "approx_kl": kl_penalty.item(), } def _compute_group_advantages( self, rewards: torch.Tensor, num_groups: int, ) -> torch.Tensor: """ 组内标准化奖励 → 优势 对每组 G 个样本，计算： A_i = (R_i - mean(R_group)) / std(R_group) 如果组内所有样本奖励相同（全对或全错），std=0，返回 0 """ # 重塑为 [num_groups, group_size] grouped_rewards = rewards.view(num_groups, self.group_size) group_mean = grouped_rewards.mean(dim=-1, keepdim=True) group_std = grouped_rewards.std(dim=-1, keepdim=True) # 避免除零：如果 std 为 0，优势全部设为 0 group_std = torch.where(group_std < 1e-8, torch.ones_like(group_std), group_std) advantages = (grouped_rewards - group_mean) / group_std # 如果原始 std 为 0，优势设为 0 raw_std = grouped_rewards.std(dim=-1, keepdim=True) advantages = torch.where(raw_std < 1e-8, torch.zeros_like(advantages), advantages) return advantages.view(-1) # 恢复为 [batch_size]

import re from typing import List, Dict, Optional class MathRewardFunction: """ 数学推理任务的奖励函数 奖励层次： 1. 答案完全正确: +1.0 2. 答案格式正确但数值错误: +0.2（鼓励正确格式） 3. 推理过程有合理步骤: +0.1（鼓励逐步推理） 4. 答案格式错误: -0.5（惩罚乱答） 5. 空答案或无关输出: -1.0 """ def __init__(self): # 用于提取答案的正则 self.answer_pattern = re.compile( r'###?\s*答案[：:]\s*(.*)', re.IGNORECASE ) self.boxed_pattern = re.compile( r'\\boxed\{([^}]*)\}' ) self.final_pattern = re.compile( r'[Tt]he\s+answer\s+is\s+[：:]*\s*([^.\n]+)', re.IGNORECASE ) def __call__( self, predictions: List[str], ground_truths: List[str], ) -> List[float]: """ 计算奖励 Args: predictions: 模型生成的回答列表 ground_truths: 正确答案列表 Returns: 奖励分数列表 """ rewards = [] for pred, gt in zip(predictions, ground_truths): reward = self._compute_single_reward(pred, gt) rewards.append(reward) return rewards def _compute_single_reward( self, prediction: str, ground_truth: str, ) -> float: """计算单个回答的奖励""" if not prediction or len(prediction.strip()) < 10: return -1.0 # 空或过短的回答 # 提取答案 extracted = self._extract_answer(prediction) if extracted is None: # 没有提取到答案格式，但可能有推理过程 # 检查是否包含"步"或"Step"等推理提示词 if any(word in prediction.lower() for word in ['步', 'step', '首先', '第一', '解', '解：']): return 0.1 # 有推理过程但答案格式不对 return -0.5 # 格式完全错误 # 答案格式正确，检查是否正确 if self._normalize_answer(extracted) == self._normalize_answer(ground_truth): return 1.0 # 答案格式正确但数值错误，检查是否有合理的推理 if any(word in prediction.lower() for word in ['步', 'step', '首先', '第一', '因此', '所以']): return 0.2 # 合理的推理但答案错了 return 0.0 # 格式对但内容错 def _extract_answer(self, text: str) -> Optional[str]: """从文本中提取答案""" # 尝试 boxed 格式（\boxed{...}） match = self.boxed_pattern.search(text) if match: return match.group(1).strip() # 尝试 "### 答案：" 格式 match = self.answer_pattern.search(text) if match: return match.group(1).strip() # 尝试 "The answer is" 格式 match = self.final_pattern.search(text) if match: return match.group(1).strip() return None @staticmethod def _normalize_answer(answer: str) -> str: """标准化答案用于比较""" # 去除空格和特殊字符 answer = answer.strip() answer = re.sub(r'[,\s]', '', answer) # 统一大小写 answer = answer.lower() # 移除前导零 answer = re.sub(r'^0+', '', answer) return answer

import torch from torch.utils.data import Dataset, DataLoader from typing import List, Dict, Optional import json class MathReasoningDataset(Dataset): """ 数学推理数据集 每个样本：{"question": "...", "answer": "..."} """ def __init__( self, data_path: str, tokenizer, max_prompt_length: int = 512, max_response_length: int = 1024, ): with open(data_path, 'r') as f: self.data = json.load(f) self.tokenizer = tokenizer self.max_prompt_length = max_prompt_length self.max_response_length = max_response_length def __len__(self): return len(self.data) def __getitem__(self, idx): item = self.data[idx] return { "question": item["question"], "answer": item["answer"], } def format_prompt(self, question: str) -> str: """格式化输入 prompt，引导模型逐步推理""" return f"""请解决下面的数学问题，并给出详细的推理步骤。 问题: {question} 请在最后用 "### 答案: [你的答案]" 的格式给出最终答案。 逐步推理过程: """ def collate_and_generate( batch: List[Dict], model, tokenizer, trainer: GRPOTrainer, reward_fn: MathRewardFunction, device: str = "cuda", ) -> Dict: """ 为 GRPO 训练准备一批数据： 1. 格式化 prompt 2. 使用旧策略采样 group_size 个输出 3. 计算奖励 4. 组装为训练 batch """ questions = [item["question"] for item in batch] answers = [item["answer"] for item in batch] all_input_ids = [] all_attention_masks = [] all_rewards = [] # 对每个问题采样 group_size 个回答 for question, gt_answer in zip(questions, answers): prompt = batch[0].get("formatted_prompt", MathReasoningDataset.format_prompt(None, question)) # 用旧策略采样 G 个输出 outputs = sample_from_model( model, tokenizer, prompt, num_samples=trainer.group_size, max_new_tokens=trainer.max_response_length, # 需要设置 temperature=0.8, # 较高温度增加多样性 top_p=0.95, ) for output_text in outputs: # 编码 prompt + 回答 full_text = prompt + output_text encoded = tokenizer( full_text, max_length=trainer.max_prompt_length + trainer.max_response_length, padding="max_length", truncation=True, return_tensors="pt", ) all_input_ids.append(encoded["input_ids"][0]) all_attention_masks.append(encoded["attention_mask"][0]) # 计算奖励 rewards = reward_fn(outputs, [gt_answer] * len(outputs)) all_rewards.extend(rewards) return { "input_ids": torch.stack(all_input_ids).to(device), "attention_mask": torch.stack(all_attention_masks).to(device), "rewards": torch.tensor(all_rewards, dtype=torch.float).to(device), } def sample_from_model( model: nn.Module, tokenizer, prompt: str, num_samples: int = 8, max_new_tokens: int = 1024, temperature: float = 0.8, top_p: float = 0.95, ) -> List[str]: """ 从模型中采样多个输出用于 GRPO 组内比较 """ inputs = tokenizer(prompt, return_tensors="pt").to(model.device) outputs = model.generate( **inputs, max_new_tokens=max_new_tokens, num_return_sequences=num_samples, do_sample=True, temperature=temperature, top_p=top_p, pad_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id, ) # 解码输出 decoded = [] prompt_len = inputs["input_ids"].shape[1] for i in range(num_samples): response = outputs[i][prompt_len:] text = tokenizer.decode(response, skip_special_tokens=True) decoded.append(text) return decoded

def train_grpo( trainer: GRPOTrainer, dataset: MathReasoningDataset, reward_fn: MathRewardFunction, num_epochs: int = 3, batch_size: int = 4, # 每批的问题数 device: str = "cuda", save_path: str = "./grpo_checkpoint", log_interval: int = 10, ): """ GRPO 训练主循环 """ dataloader = DataLoader( dataset, batch_size=batch_size, shuffle=True, collate_fn=lambda batch: batch, ) global_step = 0 all_losses = [] all_rewards = [] trainer.policy_model.train() for epoch in range(num_epochs): epoch_loss = 0.0 epoch_reward = 0.0 num_batches = 0 for step, batch in enumerate(dataloader): # ---- Step 1: 采样 ---- # 冻结旧策略，采样 group_size 个输出 trainer.policy_model.eval() with torch.no_grad(): training_data = collate_and_generate( batch, trainer.policy_model, trainer.tokenizer, trainer, reward_fn, device, ) trainer.policy_model.train() # ---- Step 2: 计算 GRPO 损失 ---- loss, stats = trainer.grpo_loss( training_data["input_ids"], training_data["attention_mask"], training_data["rewards"], ) # ---- Step 3: 反向传播 ---- loss.backward() torch.nn.utils.clip_grad_norm_( trainer.policy_model.parameters(), trainer.max_grad_norm ) trainer.optimizer.step() trainer.optimizer.zero_grad() # ---- 统计 ---- epoch_loss += loss.item() epoch_reward += stats["mean_reward"] num_batches += 1 if global_step % log_interval == 0: print( f"Epoch {epoch+1}, Step {global_step}, " f"Loss: {loss.item():.4f}, " f"Reward: {stats['mean_reward']:.4f}, " f"Advantage: {stats['mean_advantage']:.4f}, " f"KL: {stats['approx_kl']:.4f}" ) global_step += 1 avg_loss = epoch_loss / num_batches avg_reward = epoch_reward / num_batches print(f"\nEpoch {epoch+1} Summary: Avg Loss: {avg_loss:.4f}, " f"Avg Reward: {avg_reward:.4f}\n") # 保存 checkpoint save_ckpt = f"{save_path}_epoch_{epoch+1}.pt" torch.save({ "model_state_dict": trainer.policy_model.state_dict(), "optimizer_state_dict": trainer.optimizer.state_dict(), "epoch": epoch + 1, "avg_loss": avg_loss, "avg_reward": avg_reward, }, save_ckpt) print(f"Checkpoint saved to {save_ckpt}") return all_losses, all_rewards

小明有 5 个苹果，小红给了他 3 个，然后吃了 2 个。 还剩 5 + 3 - 2 = 8 - 2 = 6 个。

首先，小明原本有 5 个苹果。 然后，小红给了他 3 个，所以现在有 5 + 3 = 8 个。 接着，他吃了 2 个，所以从 8 个中减去 2 个。 最后，8 - 2 = 6 个。 让我们验证一下：5 + 3 = 8, 8 - 2 = 6。 所以还剩 6 个苹果。 ### 答案: 6

**分析**：这是一个包含加法和减法的两步问题。 1. 初始数量: 5 个苹果 2. 增加: 小红给了 3 个 → +3 3. 减少: 吃了 2 个 → -2 4. 计算: 5 + 3 = 8, 8 - 2 = 6 **检查**：5 + 3 - 2 = 6 ✓ ### 答案: 6

# KL 惩罚的自适应调整策略 class AdaptiveKLPenalty: """ 自适应调整 KL 惩罚系数 核心思想：维持 KL 散度在目标范围内 - 如果 KL 过高 → 增加惩罚 (β↑) - 如果 KL 过低 → 减少惩罚 (β↓) """ def __init__( self, initial_beta: float = 0.04, target_kl: float = 3.0, adjustment_rate: float = 0.1, ): self.beta = initial_beta self.target_kl = target_kl self.adjustment_rate = adjustment_rate def update(self, current_kl: float): if current_kl > self.target_kl * 1.5: self.beta *= (1 + self.adjustment_rate) elif current_kl < self.target_kl * 0.5: self.beta *= (1 - self.adjustment_rate) # 限制范围防止极端值 self.beta = max(0.001, min(0.5, self.beta)) return self.beta

class ShapedMathReward(MathRewardFunction): """ 带部分奖励的数学推理奖励函数 奖励拆解： - 最终答案正确: +1.0 - 中间结果正确数比例: 0~0.5 (鼓励分步正确) - 输出格式规范: 0~0.3 (鼓励结构化输出) - 推理步骤完整性: 0~0.2 (鼓励完整推理) """ def __init__(self): super().__init__() # 简单的分步模式识别 self.step_patterns = [ r'首先[，,]?\s*(.*?)[，,。]', r'然后[，,]?\s*(.*?)[，,。]', r'接着[，,]?\s*(.*?)[，,。]', r'最后[，,]?\s*(.*?)[，,。]', r'(?:第[一二三四五六七八九十]+步)[：:]\s*(.*?)[，,。]', r'Step\s*\d+[：:]\s*(.*?)[，,。\n]', ] def compute_shaped_reward( self, prediction: str, ground_truth: str, intermediate_results: Optional[List[str]] = None, ) -> float: """ 计算塑形奖励 """ base_reward = self._compute_single_reward(prediction, ground_truth) if base_reward >= 1.0: return 1.0 # 完全正确，无需额外奖励 # 格式奖励 format_score = self._compute_format_score(prediction) # 推理步骤完整性奖励 step_score = self._compute_step_score(prediction) # 总奖励 = 基础奖励 + 辅助奖励（有上限） total = base_reward + 0.3 * format_score + 0.2 * step_score return min(total, 1.0) def _compute_format_score(self, text: str) -> float: """格式规范评分""" score = 0.0 # 有明确的推理结构 if '### 答案:' in text or '\\boxed{' in text: score += 0.15 if any(p in text.lower() for p in ['首先', '然后', '接着', '最后', 'step']): score += 0.1 if '验证' in text or '检查' in text or '检验' in text: score += 0.05 return score def _compute_step_score(self, text: str) -> float: """推理步骤完整性评分""" seen_steps = 0 for pattern in self.step_patterns: if re.search(pattern, text): seen_steps += 1 # 至少 3 个步骤才算完整推理 return min(seen_steps / 5.0, 1.0)

# 使用 vLLM 加速 GRPO 采样 from vllm import LLM, SamplingParams class VLLMSampler: """基于 vLLM 的高效采样器""" def __init__(self, model_path: str, tensor_parallel_size: int = 1): self.llm = LLM( model=model_path, tensor_parallel_size=tensor_parallel_size, gpu_memory_utilization=0.8, ) self.sampling_params = SamplingParams( temperature=0.8, top_p=0.95, max_tokens=1024, ) def sample( self, prompts: List[str], num_samples: int = 8, ) -> List[List[str]]: """ 对每个 prompt 采样 num_samples 个输出 返回: [[output_1, ..., output_G], ...] """ # 为每个 prompt 复制 num_samples 份 batched_prompts = [] for prompt in prompts: batched_prompts.extend([prompt] * num_samples) outputs = self.llm.generate(batched_prompts, self.sampling_params) # 重组为 [prompts, group_size] results = [] for i in range(0, len(outputs), num_samples): group = [out.outputs[0].text for out in outputs[i:i+num_samples]] results.append(group) return results