import torch import torch.optim as optim import torch.nn.functional as F from torch.utils.data import DataLoader, TensorDataset from typing import Dict, List, Tuple, Optional from tqdm import tqdm import numpy as np import gc import logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) class GRPOTrainer: def __init__( self, actor_model, reward_model, ref_model, tokenizer, learning_rate: float = 1e-6, kl_coef: float = 0.04, group_size: int = 4, clip_epsilon: float = 0.2, max_grad_norm: float = 1.0, grpo_epochs: int = 1, update_batch_size: int = 4, use_amp: bool = True, value_clip: bool = False, entropy_coef: float = 0.01, advantage_normalization: str = 'group', # 'group', 'global', 'none' kl_estimation_method: str = 'forward' # 'forward', 'reverse', 'symmetric' ): self.actor = actor_model self.reward_model = reward_model self.ref_model = ref_model self.tokenizer = tokenizer self.kl_coef = kl_coef self.group_size = group_size self.clip_epsilon = clip_epsilon self.max_grad_norm = max_grad_norm self.grpo_epochs = grpo_epochs self.update_batch_size = update_batch_size self.use_amp = use_amp self.entropy_coef = entropy_coef self.advantage_normalization = advantage_normalization self.kl_estimation_method = kl_estimation_method self.device = next(actor_model.parameters()).device # 冻结参考模型和奖励模型 self.ref_model.eval() self.ref_model.requires_grad_(False) self.reward_model.eval() self.reward_model.requires_grad_(False) # 优化器配置 self.optimizer = optim.AdamW( filter(lambda p: p.requires_grad, actor_model.parameters()), lr=learning_rate, weight_decay=0.01, betas=(0.9, 0.95), eps=1e-8 ) # 混合精度训练 self.scaler = torch.amp.GradScaler('cuda', enabled=self.use_amp) self.training_stats = { 'iterations': 0, 'total_samples': 0, 'avg_rewards': [], 'avg_kl': [], 'policy_losses': [] } logger.info(f"GRPO Trainer initialized:") logger.info(f" Group Size: {group_size}") logger.info(f" KL Coef: {kl_coef}") logger.info(f" Clip Epsilon: {clip_epsilon}") logger.info(f" Learning Rate: {learning_rate}") logger.info(f" Update Batch Size: {update_batch_size}") logger.info(f" Mixed Precision: {use_amp}") logger.info(f" KL Estimation: {kl_estimation_method}") def _compute_kl_divergence( self, log_probs: torch.Tensor, ref_log_probs: torch.Tensor, mask: torch.Tensor ) -> torch.Tensor: if self.kl_estimation_method == 'forward': kl = log_probs - ref_log_probs elif self.kl_estimation_method == 'reverse': kl = ref_log_probs - log_probs else: forward_kl = log_probs - ref_log_probs reverse_kl = ref_log_probs - log_probs kl = 0.5 * (forward_kl + reverse_kl) kl_penalty = (kl * mask).sum(dim=-1) return kl_penalty @torch.no_grad() def generate_experience( self, prompts_dataloader: DataLoader, max_gen_len: int, temperature: float = 1.0, top_p: float = 0.9 ) -> Dict: self.actor.eval() all_sequences = [] all_log_probs = [] all_advantages = [] all_prompt_lens = [] all_rewards = [] logger.info("Generating experience...") for prompts in tqdm(prompts_dataloader, desc="Generating Experience"): try: # 处理不同的输入格式 if isinstance(prompts, (list, tuple)): prompts = prompts[0] prompts = prompts.to(self.device) batch_size = prompts.shape[0] # 扩展prompts以生成group_size个样本 prompts_repeated = prompts.repeat_interleave(self.group_size, dim=0) prompt_len = prompts_repeated.shape[1] input_data = { 'segments': [{ 'type': 'text', 'data': prompts_repeated, 'modality_id': 0 }] } # 1. 采样生成 with torch.amp.autocast('cuda', enabled=self.use_amp): response_ids = self.actor.generate( input_data, max_new_tokens=max_gen_len, do_sample=True, temperature=temperature, top_p=top_p, eos_token_id=self.tokenizer.eos_token_id, pad_token_id=self.tokenizer.pad_token_id, use_cache=True ) sequences = torch.cat([prompts_repeated, response_ids], dim=1) # 检查序列长度 if sequences.shape[1] <= prompt_len: logger.warning("Generated sequence too short, skipping batch") continue full_input_data = { 'segments': [{ 'type': 'text', 'data': sequences, 'modality_id': 0 }] } # 2. 计算当前策略和参考策略的 LogProbs with torch.amp.autocast('cuda', enabled=self.use_amp): actor_out = self.actor(full_input_data) ref_out = self.ref_model(full_input_data) logits = actor_out['logits'][:, :-1, :] ref_logits = ref_out['logits'][:, :-1, :] targets = sequences[:, 1:] log_probs = F.log_softmax(logits, dim=-1) ref_log_probs = F.log_softmax(ref_logits, dim=-1) # 提取对应token的log概率 per_token_log_probs = torch.gather( log_probs, -1, targets.unsqueeze(-1) ).squeeze(-1) per_token_ref_log_probs = torch.gather( ref_log_probs, -1, targets.unsqueeze(-1) ).squeeze(-1) # 3. 计算 KL 散度 (只针对response部分) response_mask = torch.arange( sequences.size(1) - 1, device=self.device ) >= (prompt_len - 1) response_mask = response_mask.unsqueeze(0).expand_as(per_token_log_probs) response_mask = response_mask.float() kl_penalty = self._compute_kl_divergence( per_token_log_probs, per_token_ref_log_probs, response_mask ) with torch.amp.autocast('cuda', enabled=self.use_amp): reward_output = self.reward_model(full_input_data) # reward_model返回 (batch_size, seq_len)，取最后一个位置的奖励 if reward_output.dim() == 2: raw_rewards = reward_output[:, -1] else: raw_rewards = reward_output.squeeze(-1) # 5. 组合总奖励: R_total = R_env - β * KL total_rewards = raw_rewards - self.kl_coef * kl_penalty # 6. 计算组内相对优势 rewards_grouped = total_rewards.view(batch_size, self.group_size) if self.advantage_normalization == 'group': # 组内标准化 mean_grouped = rewards_grouped.mean(dim=1, keepdim=True) std_grouped = rewards_grouped.std(dim=1, keepdim=True) + 1e-8 advantages = (rewards_grouped - mean_grouped) / std_grouped elif self.advantage_normalization == 'global': # 全局标准化 advantages = (rewards_grouped - rewards_grouped.mean()) / ( rewards_grouped.std() + 1e-8 ) else: # 'none' advantages = rewards_grouped - rewards_grouped.mean(dim=1, keepdim=True) advantages = advantages.view(-1) # 保存数据 all_sequences.append(sequences.cpu()) all_log_probs.append(per_token_log_probs.detach().cpu()) all_advantages.append(advantages.detach().cpu()) all_prompt_lens.append( torch.full((sequences.size(0),), prompt_len, dtype=torch.long) ) all_rewards.append(total_rewards.detach().cpu()) # 清理中间变量 del logits, ref_logits, actor_out, ref_out del log_probs, ref_log_probs, reward_output except Exception as e: logger.error(f"Error generating experience for batch: {e}") import traceback traceback.print_exc() continue finally: torch.cuda.empty_cache() if not all_sequences: raise RuntimeError("No valid sequences generated") # 合并所有数据 experience = { 'sequences': torch.cat(all_sequences, dim=0), 'log_probs': torch.cat(all_log_probs, dim=0), 'advantages': torch.cat(all_advantages, dim=0), 'prompt_lengths': torch.cat(all_prompt_lens, dim=0), 'rewards': torch.cat(all_rewards, dim=0) } # 统计信息 logger.info(f"Generated {len(experience['sequences'])} sequences") logger.info(f"Avg Reward: {experience['rewards'].mean().item():.4f}") logger.info(f"Reward Std: {experience['rewards'].std().item():.4f}") logger.info(f"Avg Advantage: {experience['advantages'].mean().item():.4f}") return experience def grpo_step( self, dataset: TensorDataset ) -> Dict[str, float]: self.actor.train() dataloader = DataLoader( dataset, batch_size=self.update_batch_size, shuffle=True, drop_last=False ) epoch_stats = { 'total_loss': 0.0, 'policy_loss': 0.0, 'entropy': 0.0, 'approx_kl': 0.0, 'clip_fraction': 0.0, 'steps': 0 } for batch_data in dataloader: sequences, old_log_probs, advantages, prompt_lens = batch_data sequences = sequences.to(self.device) old_log_probs = old_log_probs.to(self.device) advantages = advantages.to(self.device) input_data = { 'segments': [{ 'type': 'text', 'data': sequences, 'modality_id': 0 }] } with torch.amp.autocast('cuda', enabled=self.use_amp): outputs = self.actor(input_data) logits = outputs['logits'][:, :-1, :] # 计算新的log probabilities targets = sequences[:, 1:] log_probs_dist = F.log_softmax(logits, dim=-1) new_log_probs = torch.gather( log_probs_dist, -1, targets.unsqueeze(-1) ).squeeze(-1) # 构建response mask mask = torch.zeros_like(new_log_probs) for i, pl in enumerate(prompt_lens): mask[i, pl-1:] = 1.0 # 计算概率比率 ratio = torch.exp(new_log_probs - old_log_probs) # 扩展advantages到序列维度 adv_expanded = advantages.unsqueeze(-1).expand_as(new_log_probs) # PPO clip损失 surr1 = ratio * adv_expanded surr2 = torch.clamp( ratio, 1.0 - self.clip_epsilon, 1.0 + self.clip_epsilon ) * adv_expanded # 策略损失 policy_loss = -torch.min(surr1, surr2) policy_loss = (policy_loss * mask).sum() / (mask.sum() + 1e-8) # 熵奖励 probs = F.softmax(logits, dim=-1) entropy = -(probs * log_probs_dist).sum(dim=-1) entropy_bonus = (entropy * mask).sum() / (mask.sum() + 1e-8) # 总损失 loss = policy_loss - self.entropy_coef * entropy_bonus # 统计信息 with torch.no_grad(): log_ratio = new_log_probs - old_log_probs approx_kl = ((ratio - 1) - log_ratio) * mask approx_kl = approx_kl.sum() / (mask.sum() + 1e-8) clip_fraction = ((ratio > 1 + self.clip_epsilon) | (ratio < 1 - self.clip_epsilon)).float() clip_fraction = (clip_fraction * mask).sum() / (mask.sum() + 1e-8) self.optimizer.zero_grad() self.scaler.scale(loss).backward() # 梯度裁剪 self.scaler.unscale_(self.optimizer) grad_norm = torch.nn.utils.clip_grad_norm_( self.actor.parameters(), self.max_grad_norm ) self.scaler.step(self.optimizer) self.scaler.update() # 累积统计 epoch_stats['total_loss'] += loss.item() epoch_stats['policy_loss'] += policy_loss.item() epoch_stats['entropy'] += entropy_bonus.item() epoch_stats['approx_kl'] += approx_kl.item() epoch_stats['clip_fraction'] += clip_fraction.item() epoch_stats['steps'] += 1 # 计算平均值 for key in epoch_stats: if key != 'steps': epoch_stats[key] /= max(epoch_stats['steps'], 1) return epoch_stats def train( self, prompt_dataloader: DataLoader, num_iterations: int = 1, max_gen_len: int = 50, temperature: float = 1.0, save_every: int = 5, save_path: str = "checkpoints" ): logger.info(f"\n{'='*80}") logger.info(f"Starting GRPO Training") logger.info(f" Iterations: {num_iterations}") logger.info(f" Max Gen Length: {max_gen_len}") logger.info(f" Temperature: {temperature}") logger.info(f"{'='*80}\n") for iteration in range(num_iterations): try: # 1. 生成经验 experience = self.generate_experience( prompt_dataloader, max_gen_len, temperature ) dataset = TensorDataset( experience['sequences'], experience['log_probs'], experience['advantages'], experience['prompt_lengths'] ) # 2. 策略优化 logger.info(f"Optimizing policy for {self.grpo_epochs} epochs...") all_epoch_stats = [] for epoch in range(self.grpo_epochs): stats = self.grpo_step(dataset) all_epoch_stats.append(stats) logger.info( f" Epoch {epoch+1}/{self.grpo_epochs} | " f"Loss: {stats['total_loss']:.4f} | " f"KL: {stats['approx_kl']:.4f} | " f"Clip%: {stats['clip_fraction']*100:.1f}" ) # 3. 汇总统计 avg_stats = { key: np.mean([s[key] for s in all_epoch_stats]) for key in all_epoch_stats[0].keys() } self.training_stats['iterations'] += 1 self.training_stats['total_samples'] += len(experience['sequences']) self.training_stats['avg_rewards'].append( experience['rewards'].mean().item() ) self.training_stats['avg_kl'].append(avg_stats['approx_kl']) self.training_stats['policy_losses'].append(avg_stats['policy_loss']) # 4. 打印进度 logger.info(f"\n{'='*80}") logger.info(f"Iteration {iteration+1}/{num_iterations} Complete") logger.info(f" Avg Reward: {experience['rewards'].mean():.4f}") logger.info(f" Avg Advantage: {experience['advantages'].mean():.4f}") logger.info(f" Policy Loss: {avg_stats['policy_loss']:.4f}") logger.info(f" Approx KL: {avg_stats['approx_kl']:.4f}") logger.info(f" Entropy: {avg_stats['entropy']:.4f}") logger.info(f" Clip Fraction: {avg_stats['clip_fraction']*100:.1f}%") logger.info(f"{'='*80}\n") # 5. 保存checkpoint if (iteration + 1) % save_every == 0: self.save_checkpoint( f"{save_path}/grpo_iter_{iteration+1}.pt" ) # 6. 清理内存 del experience, dataset gc.collect() torch.cuda.empty_cache() except Exception as e: logger.error(f"Error in iteration {iteration+1}: {e}") import traceback traceback.print_exc() continue logger.info("GRPO Training Complete!") self.print_training_summary() def save_checkpoint(self, path: str): import os os.makedirs(os.path.dirname(path), exist_ok=True) checkpoint = { 'actor_state_dict': self.actor.state_dict(), 'optimizer_state_dict': self.optimizer.state_dict(), 'scaler_state_dict': self.scaler.state_dict(), 'training_stats': self.training_stats, 'config': { 'kl_coef': self.kl_coef, 'group_size': self.group_size, 'clip_epsilon': self.clip_epsilon, } } torch.save(checkpoint, path) logger.info(f"Checkpoint saved to {path}") def load_checkpoint(self, path: str): checkpoint = torch.load(path, map_location=self.device) self.actor.load_state_dict(checkpoint['actor_state_dict']) self.optimizer.load_state_dict(checkpoint['optimizer_state_dict']) if 'scaler_state_dict' in checkpoint and self.use_amp: self.scaler.load_state_dict(checkpoint['scaler_state_dict']) self.training_stats = checkpoint['training_stats'] logger.info(f"Checkpoint loaded from {path}") def print_training_summary(self): logger.info("\n" + "="*80) logger.info("Training Summary") logger.info("="*80) logger.info(f"Total Iterations: {self.training_stats['iterations']}") logger.info(f"Total Samples: {self.training_stats['total_samples']}") if self.training_stats['avg_rewards']: logger.info( f"Final Avg Reward: " f"{self.training_stats['avg_rewards'][-1]:.4f}" ) logger.info( f"Reward Improvement: " f"{self.training_stats['avg_rewards'][-1] - self.training_stats['avg_rewards'][0]:.4f}" ) logger.info("="*80 + "\n")