Update grpo.py

grpo.py

@@ -1,630 +1,539 @@
-"""
-改进的 GRPO (Group Relative Policy Optimization) 训练器
-修复了所有已知问题
-"""
-
-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:
-    """
-    GRPO训练器 - Group Relative Policy Optimization
-    参考 DeepSeekMath/DeepSeek-V3 策略
-    
-    主要修复：
-    1. 修复了 generate() 返回格式问题
-    2. 修复了 reward_model 输出处理
-    3. 添加了完整的混合精度训练支持
-    4. 改进了 KL 散度计算的数值稳定性
-    5. 修复了 past_key_values 的使用逻辑
-    6. 改进了内存管理和错误处理
-    """
-    
-    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'
-    ):
-        """
-        初始化GRPO训练器
-        
-        Args:
-            actor_model: 要训练的策略模型
-            reward_model: 奖励模型（冻结）
-            ref_model: 参考模型（冻结）
-            tokenizer: 分词器
-            learning_rate: 学习率
-            kl_coef: KL散度惩罚系数
-            group_size: 每个prompt生成的样本数
-            clip_epsilon: PPO clip范围
-            max_grad_norm: 梯度裁剪阈值
-            grpo_epochs: 每批经验的训练轮数
-            update_batch_size: 更新时的mini-batch大小
-            use_amp: 是否使用混合精度训练
-            value_clip: 是否对value进行clip（当前未使用value网络）
-            entropy_coef: 熵正则化系数
-            advantage_normalization: 优势函数归一化方式
-            kl_estimation_method: KL散度估计方法
-        """
-        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
-        )
-        
-        # 混合精度训练 - 修复：添加 GradScaler
-        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:
-        """
-        计算KL散度（改进数值稳定性）
-        
-        Args:
-            log_probs: 当前策略的log概率
-            ref_log_probs: 参考策略的log概率
-            mask: 有效token的mask
-            
-        Returns:
-            KL散度（标量）
-        """
-        if self.kl_estimation_method == 'forward':
-            # KL(π||π_ref) = Σ π * log(π/π_ref)
-            # ≈ Σ exp(log_π) * (log_π - log_π_ref)
-            # 为了数值稳定，使用 log_π - log_π_ref
-            kl = log_probs - ref_log_probs
-        elif self.kl_estimation_method == 'reverse':
-            # KL(π_ref||π) = Σ π_ref * log(π_ref/π)
-            kl = ref_log_probs - log_probs
-        else:  # symmetric
-            # 对称KL散度
-            forward_kl = log_probs - ref_log_probs
-            reverse_kl = ref_log_probs - log_probs
-            kl = 0.5 * (forward_kl + reverse_kl)
-        
-        # 应用mask并求和
-        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:
-        """
-        生成经验数据：采样 -> 计算 LogProbs -> 计算 Rewards(含KL)
-        
-        修复：
-        1. 正确处理 generate() 的返回值
-        2. 修复 reward_model 的输出处理
-        3. 改进数值稳定性
-        """
-        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. 采样生成（修复：generate只返回新生成的tokens）
-                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  # 使用缓存加速生成
-                    )
-                
-                # 修复：拼接完整序列（prompt + response）
-                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 probabilities（改进数值稳定性）
-                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计算
-                kl_penalty = self._compute_kl_divergence(
-                    per_token_log_probs,
-                    per_token_ref_log_probs,
-                    response_mask
-                )
-                
-                # 4. 计算环境奖励（修复：正确处理reward_model输出）
-                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. 计算组内相对优势 (Group Relative Advantage)
-                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]:
-        """
-        执行 GRPO 优化步骤
-        
-        修复：
-        1. 使用 GradScaler 进行混合精度训练
-        2. 改进损失计算
-        3. 更好的统计信息收集
-        """
-        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
-                }]
-            }
-            
-            # 修复：使用 GradScaler 进行混合精度训练
-            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)
-            
-            # 修复：使用 GradScaler 进行反向传播
-            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"
-    ):
-        """
-        完整的GRPO训练循环
-        
-        Args:
-            prompt_dataloader: 提供prompts的数据加载器
-            num_iterations: 训练迭代次数
-            max_gen_len: 最大生成长度
-            temperature: 采样温度
-            save_every: 每N次迭代保存一次checkpoint
-            save_path: checkpoint保存路径
-        """
-        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. 策略优化（多个epoch）
-                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):
-        """保存训练checkpoint"""
-        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(),  # 修复：保存scaler状态
-            '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"""
-        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'])
-        
-        # 修复：加载scaler状态
-        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}"
-            )
-        
+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")