Update dcpo_train.py

dcpo_train.py

@@ -17,7 +17,6 @@ warnings.filterwarnings("ignore", category=FutureWarning)
 from model import MultiModalDenseTransformer
 from dcpo import DCPOTrainer 
 
-# ================= DDP 设置 =================
 def setup_distributed():
     if "RANK" in os.environ and "WORLD_SIZE" in os.environ:
         dist.init_process_group(backend="nccl")
@@ -38,7 +37,6 @@ IS_MAIN = RANK == 0
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO if IS_MAIN else logging.WARNING)
 
-# ================= 数据集 =================
 class MathDataset(Dataset):
     def __init__(self, path):
         self.data = []
@@ -59,9 +57,7 @@ def math_collate(batch):
         'ground_truth': [item['ground_truth'] for item in batch]
     }
 
-# ================= 主函数 =================
 def main():
-    # ------------------ 配置区域 ------------------
     CONFIG = {
         'sft_checkpoint': '/root/checkpoints/dcpo_posttrain_round3/step_1200.pt', 
         'data_path': '/root/dataset/r1_zero_math.jsonl',
@@ -87,15 +83,12 @@ def main():
         'gradient_accumulation_steps': 8, 
         'inner_batch_size': 4,
         
-        # ========== 关键新增1: 奖励验证器配置 ==========
-        'use_reference_comparison': True,   # 是否使用参考推理对比
-        'use_progressive_reward': False,    # 是否使用渐进式奖励（实验性）
-        'phase1_steps': 2000,              # 渐进式阶段1（宽松格式）
-        'phase2_steps': 4000,              # 渐进式阶段2（中等格式）
+        'use_reference_comparison': True,   
+        'use_progressive_reward': False,    
+        'phase1_steps': 2000,              
+        'phase2_steps': 4000,              
     }
-    # ---------------------------------------------
     
-    # 初始化日志文件 Handler
     file_handler = None
     if IS_MAIN:
         os.makedirs(CONFIG['save_dir'], exist_ok=True)
@@ -105,31 +98,17 @@ def main():
         file_handler = logging.FileHandler(log_file, encoding='utf-8')
         file_handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))
         logger.addHandler(file_handler)
-        
-        # 将配置写入日志文件
-        logger.info(f"DCPO Configuration: {json.dumps(CONFIG, indent=2)}")
-        
-        # ========== 关键新增2: 记录使用的验证器类型 ==========
-        if CONFIG['use_progressive_reward']:
-            logger.info(f"使用渐进式奖励验证器:")
-            logger.info(f"  - 阶段1 (0-{CONFIG['phase1_steps']}): 宽松格式")
-            logger.info(f"  - 阶段2 ({CONFIG['phase1_steps']}-{CONFIG['phase2_steps']}): 中等格式")
-            logger.info(f"  - 阶段3 ({CONFIG['phase2_steps']}+): 完整要求")
-        else:
-            logger.info(f"使用标准改进版验证器 (reference_comparison={CONFIG['use_reference_comparison']})")
-        
+
         metrics_file = os.path.join(CONFIG['save_dir'], "metrics.jsonl")
         if not os.path.exists(metrics_file):
             with open(metrics_file, 'w', encoding='utf-8') as f:
                 pass 
-    
-    # 1. 加载 Tokenizer
+
     tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B-Instruct", trust_remote_code=True)
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
         tokenizer.pad_token_id = tokenizer.eos_token_id
-        
-    # 2. 初始化模型
+
     def create_model():
         return MultiModalDenseTransformer(
             model_dim=CONFIG['model_dim'],
@@ -152,8 +131,6 @@ def main():
     if WORLD_SIZE > 1:
         actor = DDP(actor, device_ids=[LOCAL_RANK], output_device=LOCAL_RANK)
 
-    # 3. 初始化 Trainer
-    # ========== 关键新增3: 传入新的验证器参数 ==========
     trainer = DCPOTrainer(
         actor_model=actor,
         ref_model=ref,
@@ -166,14 +143,12 @@ def main():
         use_amp=True,
         gradient_accumulation_steps=CONFIG['gradient_accumulation_steps'],
         inner_batch_size=CONFIG['inner_batch_size'],
-        # 新增参数
         use_reference_comparison=CONFIG['use_reference_comparison'],
         use_progressive_reward=CONFIG['use_progressive_reward'],
         phase1_steps=CONFIG['phase1_steps'],
         phase2_steps=CONFIG['phase2_steps']
     )
-    
-    # 4. 恢复状态
+
     start_step = 0
     samples_seen = 0
     
@@ -203,7 +178,6 @@ def main():
         start_step = checkpoint.get('step', 0) + 1
         samples_seen = checkpoint.get('samples_seen', start_step * CONFIG['batch_size'] * WORLD_SIZE)
         
-        # ========== 关键新增4: 恢复时更新步数（用于渐进式奖励） ==========
         if CONFIG['use_progressive_reward']:
             trainer.update_step(start_step)
             if IS_MAIN:
@@ -228,7 +202,6 @@ def main():
         gc.collect()
         torch.cuda.empty_cache()
     
-    # 5. Dataloader
     dataset = MathDataset(CONFIG['data_path'])
     if WORLD_SIZE > 1:
         sampler = torch.utils.data.DistributedSampler(
@@ -245,7 +218,6 @@ def main():
     if IS_MAIN:
         print(f"Starting Training from step {start_step}")
     
-    # 6. Dataloader 指针恢复
     if sampler:
         epoch = start_step // len(dataloader)
         sampler.set_epoch(epoch)
@@ -267,7 +239,6 @@ def main():
                 data_iter = iter(dataloader)
                 next(data_iter)
             
-    # 7. 训练循环
     progress_bar = tqdm(
         range(start_step, CONFIG['max_steps']), 
         disable=not IS_MAIN, 
@@ -282,7 +253,6 @@ def main():
     
     for step in progress_bar:
         try:
-            # ========== 关键新增5: 更新训练步数（用于渐进式奖励） ==========
             if CONFIG['use_progressive_reward']:
                 trainer.update_step(step)
             
@@ -297,7 +267,7 @@ def main():
             
             samples_seen += CONFIG['batch_size'] * WORLD_SIZE
             
-            # 1. 生成 + SAS
+            # 生成 + SAS
             experience = trainer.generate_and_prepare(
                 batch, 
                 max_gen_len=CONFIG['max_gen_len']
@@ -307,13 +277,11 @@ def main():
             if running_reward == 0: running_reward = step_reward
             else: running_reward = 0.95 * running_reward + 0.05 * step_reward
             
-            # 2. 训练步骤
+
             loss = trainer.train_step(experience)
             
-            # 状态栏缩写
             status_dict = {"Rw": f"{running_reward:.2f}"}
-            
-            # ========== 关键新增6: 添加阶段信息显示（如果使用渐进式） ==========
+
             if CONFIG['use_progressive_reward'] and hasattr(trainer.math_verifier, 'current_phase'):
                 status_dict["Ph"] = f"{trainer.math_verifier.current_phase}"
             
@@ -334,7 +302,6 @@ def main():
                         "timestamp": datetime.now().isoformat()
                     }
                     
-                    # ========== 关键新增7: 记录渐进式阶段信息 ==========
                     if CONFIG['use_progressive_reward'] and hasattr(trainer.math_verifier, 'current_phase'):
                         metrics_data['reward_phase'] = trainer.math_verifier.current_phase
                     
@@ -353,8 +320,7 @@ def main():
                 status_dict["St"] = "Acc"
             
             progress_bar.set_description(f"{' '.join([f'{k}:{v}' for k,v in status_dict.items()])}")
-            
-            # 保存逻辑
+
             is_accum_boundary = (len(trainer.experience_buffer) == 0)
             
             if step > 0 and step % CONFIG['save_interval'] == 0 and IS_MAIN: