Update dcpo.py

dcpo.py

@@ -5,19 +5,11 @@ import numpy as np
 import logging
 import hashlib
 
-# ========== 关键修改1: 导入改进的验证器 ==========
 from math_verifier import MathReward
-# 如果要使用渐进式奖励，取消下面的注释：
-# from progressive_reward import ProgressiveMathReward
 
 logger = logging.getLogger(__name__)
 
 class DCPOTrainer:
-    """
-    DCPO: Dynamic Clipping Policy Optimization Trainer
-    修复版：包含 DDP 设备修复、优化器状态恢复修复、显存优化、attention_mask 和 position_ids 修复。
-    改进版：集成改进的奖励验证器
-    """
     def __init__(
         self,
         actor_model,
@@ -33,21 +25,18 @@ class DCPOTrainer:
         use_amp: bool = True,
         gradient_accumulation_steps: int = 1,
         inner_batch_size: int = 4,
-        # ========== 关键修改2: 新增参数 ==========
-        use_reference_comparison: bool = True,  # 是否使用参考推理对比
-        use_progressive_reward: bool = False,   # 是否使用渐进式奖励
-        phase1_steps: int = 2000,              # 渐进式阶段1步数
-        phase2_steps: int = 4000               # 渐进式阶段2步数
+        use_reference_comparison: bool = True,  
+        use_progressive_reward: bool = False,   
+        phase1_steps: int = 2000,              
+        phase2_steps: int = 4000               
     ):
         self.actor = actor_model
         self.ref_model = ref_model
         self.tokenizer = tokenizer
         
-        # ========== 关键修改3: 初始化验证器 ==========
         self.use_progressive_reward = use_progressive_reward
         
         if use_progressive_reward:
-            # 使用渐进式奖励（实验性）
             from progressive_reward import ProgressiveMathReward
             self.math_verifier = ProgressiveMathReward(
                 use_reference_comparison=use_reference_comparison,
@@ -55,13 +44,10 @@ class DCPOTrainer:
                 phase2_steps=phase2_steps,
                 verbose=True
             )
-            logger.info("使用渐进式奖励验证器")
         else:
-            # 使用标准改进版验证器（推荐）
             self.math_verifier = MathReward(
                 use_reference_comparison=use_reference_comparison
             )
-            logger.info(f"使用改进版奖励验证器 (reference_comparison={use_reference_comparison})")
         
         self.group_size = group_size
         self.eps_low = eps_low
@@ -74,51 +60,41 @@ class DCPOTrainer:
         self.gradient_accumulation_steps = gradient_accumulation_steps
         self.inner_batch_size = inner_batch_size
         self.experience_buffer = [] 
-        
-        # ========== 关键修改4: 添加当前步数跟踪（用于渐进式奖励） ==========
         self.current_step = 0
-        
-        # 自动获取设备：兼容 DDP
+
         if hasattr(actor_model, 'module'):
             self.device = next(actor_model.module.parameters()).device
         else:
             self.device = next(actor_model.parameters()).device
-        
-        # 优化器初始化
+
         self.optimizer = torch.optim.AdamW(
             self.actor.parameters(), 
             lr=learning_rate,
             weight_decay=0.01
         )
-        
-        # 混合精度 Scaler
+
         self.scaler = torch.amp.GradScaler('cuda', enabled=use_amp)
 
         if self.ref_model:
             self.ref_model.eval()
             self.ref_model.requires_grad_(False)
 
-        # SAS 统计缓存
         self.sas_stats = {}
 
     def _get_stable_hash(self, text):
-        """生成跨进程/跨运行一致的哈希值"""
         return hashlib.md5(text.encode('utf-8')).hexdigest()
 
     def state_dict(self):
-        """导出 Trainer 状态"""
         return {
             'optimizer_state_dict': self.optimizer.state_dict(),
             'sas_stats': self.sas_stats,
             'scaler_state_dict': self.scaler.state_dict() if self.scaler is not None else None,
-            'current_step': self.current_step  # 保存当前步数
+            'current_step': self.current_step  
         }
 
     def load_state_dict(self, state_dict):
-        """加载 Trainer 状态，并修复优化器 Tensor 设备问题"""
         if 'optimizer_state_dict' in state_dict:
             self.optimizer.load_state_dict(state_dict['optimizer_state_dict'])
-            # 强制将优化器状态移动到当前 GPU
             for state in self.optimizer.state.values():
                 for k, v in state.items():
                     if isinstance(v, torch.Tensor):
@@ -126,24 +102,20 @@ class DCPOTrainer:
                         
         if 'sas_stats' in state_dict:
             self.sas_stats = state_dict['sas_stats']
-            logger.info(f"Loaded SAS stats for {len(self.sas_stats)} prompts")
+
             
         if 'scaler_state_dict' in state_dict and state_dict['scaler_state_dict'] is not None:
             self.scaler.load_state_dict(state_dict['scaler_state_dict'])
         
         if 'current_step' in state_dict:
             self.current_step = state_dict['current_step']
-            logger.info(f"Loaded current_step: {self.current_step}")
 
-    # ========== 关键修改5: 新增方法用于更新步数 ==========
     def update_step(self, step):
-        """更新当前训练步数（用于渐进式奖励）"""
         self.current_step = step
         if self.use_progressive_reward:
             self.math_verifier.update_step(step)
 
     def _get_unwrapped_model(self, model):
-        """辅助函数：获取原始模型（剥离 DDP wrapper）"""
         if hasattr(model, 'module'): 
             return model.module
         return model
@@ -199,17 +171,13 @@ class DCPOTrainer:
         expanded_gts = []
         for gt in ground_truths:
             expanded_gts.extend([gt] * self.group_size)
-            
-        # ========== 计算奖励（使用改进的验证器）==========
-        # 改进的验证器会自动处理 reasoning 和 reference_completion 字段
+
         raw_rewards = self.math_verifier.compute_rewards(full_responses_for_reward, expanded_gts)
         rewards_tensor = torch.tensor(raw_rewards, device=self.device, dtype=torch.float32)
         
-        # 计算旧策略的 Log Probs
         gen_mask = (generated_ids != self.tokenizer.pad_token_id).long()
         full_attention_mask = torch.cat([attention_mask_repeated, gen_mask], dim=1)
         
-        # ✅ 修复：构建正确的 position_ids（考虑左 padding）
         batch_size = sequences.size(0)
         seq_len = sequences.size(1)
         position_ids = torch.zeros((batch_size, seq_len), dtype=torch.long, device=self.device)
@@ -229,7 +197,7 @@ class DCPOTrainer:
             actor_out = self.actor(
                 full_input_data, 
                 attention_mask=full_attention_mask,
-                position_ids=position_ids  # ✅ 添加 position_ids
+                position_ids=position_ids  
             )
         
         logits = actor_out['logits'][:, :-1, :]
@@ -237,14 +205,13 @@ class DCPOTrainer:
         log_probs = F.log_softmax(logits, dim=-1)
         per_token_log_probs = torch.gather(log_probs, -1, targets.unsqueeze(-1)).squeeze(-1)
         
-        # ✅ 显存优化：提前移到 CPU
         return {
             'prompts_text': prompts_text,
-            'sequences': sequences.detach().cpu(),  # ✅ 移到 CPU
-            'old_log_probs': per_token_log_probs.detach().cpu(),  # ✅ 移到 CPU
-            'rewards': rewards_tensor.cpu(),  # ✅ 移到 CPU
-            'attention_mask': full_attention_mask.cpu(),  # ✅ 新增：保存 attention_mask
-            'position_ids': position_ids.cpu(),  # ✅ 新增：保存 position_ids
+            'sequences': sequences.detach().cpu(),  
+            'old_log_probs': per_token_log_probs.detach().cpu(), 
+            'rewards': rewards_tensor.cpu(), 
+            'attention_mask': full_attention_mask.cpu(),  
+            'position_ids': position_ids.cpu(),  
             'prompt_length': prompt_len
         }
 
@@ -268,7 +235,6 @@ class DCPOTrainer:
         mu_old = stats['mu_total']
         var_old = stats['var_total']
         
-        # 增量更新公式
         mu_total = (mu_new + (i - 1) * mu_old) / i
         term3 = ((i - 1) / i) * (mu_old - mu_new)**2
         var_total = (var_new + (i - 1) * var_old + term3) / i
@@ -304,7 +270,6 @@ class DCPOTrainer:
         return torch.cat(final_advantages)
 
     def train_step(self, experience):
-        """执行训练步骤：梯度累积 -> PPO/GRPO Update"""
         self.experience_buffer.append(experience)
         if len(self.experience_buffer) < self.gradient_accumulation_steps:
             return None
@@ -312,7 +277,7 @@ class DCPOTrainer:
         all_advantages = []
         for exp in self.experience_buffer:
             adv = self._compute_sas_advantages(exp)
-            exp['advantages'] = adv.detach()  # 保持在 CPU
+            exp['advantages'] = adv.detach()  
             all_advantages.append(exp['advantages'])
         
         self.actor.train()
@@ -325,17 +290,15 @@ class DCPOTrainer:
 
         padded_seqs = []
         padded_old_lp = []
-        padded_attn_masks = []  # ✅ 新增
-        padded_pos_ids = []  # ✅ 新增
+        padded_attn_masks = []  
+        padded_pos_ids = []  
         prompt_lens_list = []
         
         for e in self.experience_buffer:
             padded_seqs.append(pad_tensor(e['sequences'], max_seq_len, self.tokenizer.pad_token_id))
-            
-            # ✅ 修复：使用 0.0 填充（exp(0)=1，数值稳定）
+
             padded_old_lp.append(pad_tensor(e['old_log_probs'], max_lp_len, 0.0))
             
-            # ✅ 新增：padding attention_mask 和 position_ids
             padded_attn_masks.append(pad_tensor(e['attention_mask'], max_seq_len, 0))
             padded_pos_ids.append(pad_tensor(e['position_ids'], max_seq_len, 0))
             
@@ -346,8 +309,8 @@ class DCPOTrainer:
         cat_old_log_probs = torch.cat(padded_old_lp, dim=0)
         cat_advantages = torch.cat(all_advantages, dim=0)
         cat_prompt_lens = torch.tensor(prompt_lens_list)
-        cat_attention_masks = torch.cat(padded_attn_masks, dim=0)  # ✅ 新增
-        cat_position_ids = torch.cat(padded_pos_ids, dim=0)  # ✅ 新增
+        cat_attention_masks = torch.cat(padded_attn_masks, dim=0)  
+        cat_position_ids = torch.cat(padded_pos_ids, dim=0)  
         
         self.experience_buffer = []
         
@@ -356,8 +319,8 @@ class DCPOTrainer:
             cat_old_log_probs, 
             cat_advantages, 
             cat_prompt_lens,
-            cat_attention_masks,  # ✅ 新增
-            cat_position_ids  # ✅ 新增
+            cat_attention_masks,  
+            cat_position_ids  
         )
         dataloader = DataLoader(dataset, batch_size=self.inner_batch_size, shuffle=True)
         
@@ -366,13 +329,11 @@ class DCPOTrainer:
         
         for _ in range(self.grpo_epochs):
             for batch in dataloader:
-                # ✅ 解包所有数据
                 seqs, old_lp, advs, p_lens, attn_masks, pos_ids = [b.to(self.device) for b in batch]
                 
                 input_data = {'segments': [{'type': 'text', 'data': seqs, 'modality_id': 0}]}
                 
                 with torch.amp.autocast('cuda', enabled=self.use_amp):
-                    # ✅ 修复：传入 attention_mask 和 position_ids
                     outputs = self.actor(
                         input_data, 
                         attention_mask=attn_masks,
@@ -383,22 +344,19 @@ class DCPOTrainer:
                     
                     new_log_probs = F.log_softmax(logits, dim=-1)
                     new_token_log_probs = torch.gather(new_log_probs, -1, targets.unsqueeze(-1)).squeeze(-1)
-                    
-                    # Mask 构建（保持原有逻辑）
+
                     mask = torch.zeros_like(new_token_log_probs)
                     for i, pl in enumerate(p_lens):
                         pl_val = int(pl.item())
                         start_idx = max(0, pl_val - 1)
                         if start_idx < mask.size(1):
                             mask[i, start_idx:] = 1.0
-                    
-                    # 过滤 padding 和无效的 old_log_probs
+
                     is_padding = (targets == self.tokenizer.pad_token_id)
-                    is_valid_old_lp = (old_lp != 0.0)  # ✅ 修改：过滤填充值
+                    is_valid_old_lp = (old_lp != 0.0)  
                     mask = mask * (~is_padding).float() * is_valid_old_lp.float()
                     
-                    # ✅ 修复：DCPO Loss 计算 - 数值稳定性
-                    q_probs = torch.exp(old_lp).clamp(min=1e-10, max=1.0)  # ✅ clamp 避免除零
+                    q_probs = torch.exp(old_lp).clamp(min=1e-10, max=1.0)  
                     term_low = 1.0 - (4.0 * self.eps_low) / q_probs
                     lower_clip = 0.5 + 0.5 * torch.sqrt(torch.clamp(term_low, min=0.0))
                     term_high = 1.0 + (4.0 * self.eps_high) / q_probs