Upload 2 files

dcpo.py

@@ -0,0 +1,431 @@
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, TensorDataset
+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,
+        ref_model,
+        tokenizer,
+        learning_rate: float = 1e-6,
+        group_size: int = 4,
+        eps_low: float = 0.16,  
+        eps_high: float = 0.2,
+        r_max: float = 10.0,
+        grpo_epochs: int = 1,
+        max_grad_norm: float = 1.0,
+        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步数
+    ):
+        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,
+                phase1_steps=phase1_steps,
+                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
+        self.eps_high = eps_high
+        self.r_max = r_max
+        self.grpo_epochs = grpo_epochs
+        self.use_amp = use_amp
+        self.max_grad_norm = max_grad_norm
+        
+        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  # 保存当前步数
+        }
+
+    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):
+                        state[k] = v.to(self.device)
+                        
+        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
+
+    @torch.no_grad()
+    def generate_and_prepare(self, prompt_batch, max_gen_len=512, temperature=1.0):
+        self.actor.eval()
+        prompts_text = prompt_batch['prompt']
+        ground_truths = prompt_batch['ground_truth']
+        
+        inputs = self.tokenizer(
+            prompts_text, 
+            return_tensors="pt", 
+            padding=True, 
+            padding_side="left"
+        ).to(self.device)
+        
+        prompts_ids = inputs['input_ids']
+        attention_mask = inputs['attention_mask']
+        prompt_len = int(prompts_ids.shape[1])
+        
+        prompts_ids_repeated = prompts_ids.repeat_interleave(self.group_size, dim=0)
+        attention_mask_repeated = attention_mask.repeat_interleave(self.group_size, dim=0)
+        
+        input_data = {
+            'segments': [{'type': 'text', 'data': prompts_ids_repeated, 'modality_id': 0}],
+            'attention_mask': attention_mask_repeated
+        }
+        
+        # 推理时使用 unwrapped model
+        unwrapped_actor = self._get_unwrapped_model(self.actor)
+        with torch.amp.autocast('cuda', enabled=self.use_amp):
+            generated_ids = unwrapped_actor.generate(
+                input_data,
+                max_new_tokens=max_gen_len,
+                do_sample=True,
+                temperature=temperature,
+                top_p=0.95,
+                pad_token_id=self.tokenizer.pad_token_id
+            )
+            
+        sequences = torch.cat([prompts_ids_repeated, generated_ids], dim=1)
+        decoded_responses = self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
+        
+        # 处理 Think 标签
+        full_responses_for_reward = []
+        for r in decoded_responses:
+            if not r.strip().startswith("<think>"):
+                full_responses_for_reward.append("<think>\n" + r.strip())
+            else:
+                full_responses_for_reward.append(r)
+                
+        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)
+        
+        for i in range(batch_size):
+            # 找到第一个非 padding token 的位置
+            non_pad_positions = (full_attention_mask[i] == 1).nonzero(as_tuple=True)[0]
+            if len(non_pad_positions) > 0:
+                start_pos = non_pad_positions[0].item()
+                valid_len = len(non_pad_positions)
+                # 从 0 开始编号有效 token 的位置
+                position_ids[i, start_pos:start_pos + valid_len] = torch.arange(valid_len, device=self.device)
+        
+        full_input_data = {'segments': [{'type': 'text', 'data': sequences, 'modality_id': 0}]}
+        
+        with torch.amp.autocast('cuda', enabled=self.use_amp):
+            actor_out = self.actor(
+                full_input_data, 
+                attention_mask=full_attention_mask,
+                position_ids=position_ids  # ✅ 添加 position_ids
+            )
+        
+        logits = actor_out['logits'][:, :-1, :]
+        targets = sequences[:, 1:]
+        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
+            'prompt_length': prompt_len
+        }
+
+    def _update_sas_stats(self, prompt_text, new_rewards):
+        """更新 SAS 均值和方差统计"""
+        prompt_hash = self._get_stable_hash(prompt_text)
+        
+        mu_new = new_rewards.mean().item()
+        var_new = new_rewards.var(unbiased=False).item() if len(new_rewards) > 1 else 0.0
+        
+        if prompt_hash not in self.sas_stats:
+            self.sas_stats[prompt_hash] = {
+                'i': 1,
+                'mu_total': mu_new,
+                'var_total': var_new
+            }
+            return mu_new, np.sqrt(var_new + 1e-8), mu_new, np.sqrt(var_new + 1e-8)
+        
+        stats = self.sas_stats[prompt_hash]
+        i = stats['i'] + 1
+        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
+        
+        stats['i'] = i
+        stats['mu_total'] = mu_total
+        stats['var_total'] = var_total
+        
+        return mu_new, np.sqrt(var_new + 1e-8), mu_total, np.sqrt(var_total + 1e-8)
+
+    def _compute_sas_advantages(self, experience_batch):
+        prompts = experience_batch['prompts_text']
+        rewards = experience_batch['rewards'].view(-1, self.group_size)
+        
+        final_advantages = []
+        
+        for idx, prompt in enumerate(prompts):
+            group_rewards = rewards[idx]
+            mu_new, std_new, mu_total, std_total = self._update_sas_stats(prompt, group_rewards)
+            
+            A_new = (group_rewards - mu_new) / (std_new + 1e-8)
+            A_total = (group_rewards - mu_total) / (std_total + 1e-8)
+            
+            i = self.sas_stats[self._get_stable_hash(prompt)]['i']
+            
+            SA_new = ((i - 1) / i) * A_new + (1 / i) * A_total
+            SA_total = (1 / i) * A_new + ((i - 1) / i) * A_total
+            
+            mask = (torch.abs(SA_new) < torch.abs(SA_total)).float()
+            A_final = mask * SA_new + (1 - mask) * SA_total
+            final_advantages.append(A_final)
+            
+        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
+
+        all_advantages = []
+        for exp in self.experience_buffer:
+            adv = self._compute_sas_advantages(exp)
+            exp['advantages'] = adv.detach()  # 保持在 CPU
+            all_advantages.append(exp['advantages'])
+        
+        self.actor.train()
+        
+        max_seq_len = max([e['sequences'].size(1) for e in self.experience_buffer])
+        max_lp_len = max([e['old_log_probs'].size(1) for e in self.experience_buffer])
+        
+        def pad_tensor(t, target_len, val):
+            return F.pad(t, (0, target_len - t.size(1)), value=val)
+
+        padded_seqs = []
+        padded_old_lp = []
+        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))
+            
+            prompt_lens_list.extend([e['prompt_length']] * (len(e['sequences'])))
+
+        # 显存优化：Dataset 保持在 CPU
+        cat_sequences = torch.cat(padded_seqs, dim=0)
+        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)  # ✅ 新增
+        
+        self.experience_buffer = []
+        
+        dataset = TensorDataset(
+            cat_sequences, 
+            cat_old_log_probs, 
+            cat_advantages, 
+            cat_prompt_lens,
+            cat_attention_masks,  # ✅ 新增
+            cat_position_ids  # ✅ 新增
+        )
+        dataloader = DataLoader(dataset, batch_size=self.inner_batch_size, shuffle=True)
+        
+        total_loss = 0
+        update_steps = 0
+        
+        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,
+                        position_ids=pos_ids
+                    )
+                    logits = outputs['logits'][:, :-1, :]
+                    targets = seqs[:, 1:]
+                    
+                    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)  # ✅ 修改：过滤填充值
+                    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 避免除零
+                    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
+                    upper_clip = 0.5 + 0.5 * torch.sqrt(torch.clamp(term_high, min=0.0))
+                    
+                    ratio = torch.exp(new_token_log_probs - old_lp)
+                    ratio = torch.clamp(ratio, 0, self.r_max)
+                    
+                    advs_expanded = advs.unsqueeze(1).expand_as(ratio)
+                    surr1 = ratio * advs_expanded
+                    clipped_ratio = torch.min(torch.max(ratio, lower_clip), upper_clip)
+                    surr2 = clipped_ratio * advs_expanded
+                    
+                    element_wise_loss = torch.min(surr1, surr2)
+                    masked_loss = element_wise_loss * mask
+                    response_lens = torch.clamp(mask.sum(dim=1), min=1.0)
+                    per_response_loss = masked_loss.sum(dim=1) / response_lens
+                    loss = -per_response_loss.mean()
+                
+                self.optimizer.zero_grad()
+                self.scaler.scale(loss).backward()
+                self.scaler.unscale_(self.optimizer)
+                torch.nn.utils.clip_grad_norm_(self.actor.parameters(), self.max_grad_norm)
+                self.scaler.step(self.optimizer)
+                self.scaler.update()
+                
+                total_loss += loss.item()
+                update_steps += 1
+                
+        return total_loss / max(update_steps, 1)

dcpo_train.py

@@ -0,0 +1,404 @@
+import os
+import torch
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel as DDP
+from transformers import AutoTokenizer
+from torch.utils.data import DataLoader, Dataset
+import json
+import logging
+from tqdm import tqdm
+import glob
+from datetime import datetime
+import gc
+import warnings
+
+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")
+        rank = int(os.environ["RANK"])
+        local_rank = int(os.environ["LOCAL_RANK"])
+        world_size = int(os.environ["WORLD_SIZE"])
+        torch.cuda.set_device(local_rank)
+        if rank == 0:
+            print(f"Initialized DDP: Rank {rank}/{world_size}")
+        return rank, local_rank, world_size
+    else:
+        print("Initialized Single GPU Mode")
+        return 0, 0, 1
+
+RANK, LOCAL_RANK, WORLD_SIZE = setup_distributed()
+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 = []
+        with open(path, 'r', encoding='utf-8') as f:
+            for line in f:
+                if line.strip():
+                    self.data.append(json.loads(line))
+    
+    def __len__(self):
+        return len(self.data)
+    
+    def __getitem__(self, idx):
+        return self.data[idx]
+
+def math_collate(batch):
+    return {
+        'prompt': [item['prompt'] for item in 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',
+        'save_dir': '/root/checkpoints/dcpo_training',
+        'resume_from': None,
+        
+        'model_dim': 1536,
+        'n_layers': 12,
+        'n_heads': 12,
+        'n_kv_heads': 4,
+        
+        'group_size': 4,       
+        'batch_size': 1,       
+        'learning_rate': 1e-6,
+        'max_steps': 5000,    
+        'max_gen_len': 512,
+        'save_interval': 1400,
+        
+        'dcpo_eps_low': 0.16,
+        'dcpo_eps_high': 0.2,
+        'dcpo_r_max': 10.0,
+        
+        '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（中等格式）
+    }
+    # ---------------------------------------------
+    
+    # 初始化日志文件 Handler
+    file_handler = None
+    if IS_MAIN:
+        os.makedirs(CONFIG['save_dir'], exist_ok=True)
+        current_time = datetime.now().strftime('%Y%m%d_%H%M%S')
+        log_file = os.path.join(CONFIG['save_dir'], f"dcpo_train_{current_time}.log")
+        
+        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'],
+            vocab_size=len(tokenizer),
+            n_layers=CONFIG['n_layers'],
+            n_heads=CONFIG['n_heads'],
+            n_kv_heads=CONFIG['n_kv_heads'],
+            max_seq_len=2048,
+            use_gradient_checkpointing=True
+        )
+        
+    device = torch.device(f"cuda:{LOCAL_RANK}")
+    
+    if IS_MAIN:
+        print("Initializing Actor Model...")
+    
+    actor = create_model().to(device)
+    ref = None
+
+    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,
+        tokenizer=tokenizer,
+        learning_rate=CONFIG['learning_rate'],
+        group_size=CONFIG['group_size'],
+        eps_low=CONFIG['dcpo_eps_low'],
+        eps_high=CONFIG['dcpo_eps_high'],
+        r_max=CONFIG['dcpo_r_max'],
+        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
+    
+    if CONFIG['resume_from']:
+        resume_path = CONFIG['resume_from']
+        if IS_MAIN:
+            print(f"Resuming from: {resume_path}")
+            
+        checkpoint = torch.load(resume_path, map_location='cpu')
+        
+        if WORLD_SIZE > 1:
+            actor.module.load_state_dict(checkpoint['model_state_dict'])
+        else:
+            actor.load_state_dict(checkpoint['model_state_dict'])
+            
+        if 'trainer_state_dict' in checkpoint:
+            trainer.load_state_dict(checkpoint['trainer_state_dict'])
+            
+        if 'rng_state' in checkpoint:
+            torch.set_rng_state(checkpoint['rng_state'])
+        if 'cuda_rng_state' in checkpoint:
+            try:
+                torch.cuda.set_rng_state_all(checkpoint['cuda_rng_state'])
+            except:
+                torch.cuda.set_rng_state(checkpoint['cuda_rng_state'][LOCAL_RANK])
+        
+        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:
+                print(f"Restored progressive reward state to step {start_step}")
+        
+        del checkpoint
+        gc.collect()
+        torch.cuda.empty_cache() 
+    else:
+        if IS_MAIN:
+            print(f"Loading SFT checkpoint: {CONFIG['sft_checkpoint']}")
+        checkpoint = torch.load(CONFIG['sft_checkpoint'], map_location='cpu')
+        state_dict = checkpoint['model_state_dict'] if 'model_state_dict' in checkpoint else checkpoint
+        new_state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()}
+        
+        if WORLD_SIZE > 1:
+            actor.module.load_state_dict(new_state_dict)
+        else:
+            actor.load_state_dict(new_state_dict)
+            
+        del checkpoint, state_dict, new_state_dict
+        gc.collect()
+        torch.cuda.empty_cache()
+    
+    # 5. Dataloader
+    dataset = MathDataset(CONFIG['data_path'])
+    if WORLD_SIZE > 1:
+        sampler = torch.utils.data.DistributedSampler(
+            dataset, num_replicas=WORLD_SIZE, rank=RANK, shuffle=True, seed=42
+        )
+    else:
+        sampler = None
+    
+    dataloader = DataLoader(
+        dataset, batch_size=CONFIG['batch_size'], 
+        collate_fn=math_collate, sampler=sampler, shuffle=(sampler is None)
+    )
+    
+    if IS_MAIN:
+        print(f"Starting Training from step {start_step}")
+    
+    # 6. Dataloader 指针恢复
+    if sampler:
+        epoch = start_step // len(dataloader)
+        sampler.set_epoch(epoch)
+    
+    data_iter = iter(dataloader)
+    steps_in_epoch = start_step % len(dataloader)
+    
+    if start_step > 0 and steps_in_epoch > 0:
+        if IS_MAIN:
+            print(f"Fast-forwarding dataloader by {steps_in_epoch} steps...")
+        
+        for _ in range(steps_in_epoch):
+            try:
+                next(data_iter)
+            except StopIteration:
+                if sampler: 
+                    epoch += 1
+                    sampler.set_epoch(epoch)
+                data_iter = iter(dataloader)
+                next(data_iter)
+            
+    # 7. 训练循环
+    progress_bar = tqdm(
+        range(start_step, CONFIG['max_steps']), 
+        disable=not IS_MAIN, 
+        initial=start_step, 
+        total=CONFIG['max_steps'],
+        ncols=120,
+        mininterval=1.0 
+    )
+    
+    running_reward = 0.0
+    running_loss = 0.0
+    
+    for step in progress_bar:
+        try:
+            # ========== 关键新增5: 更新训练步数（用于渐进式奖励） ==========
+            if CONFIG['use_progressive_reward']:
+                trainer.update_step(step)
+            
+            try:
+                batch = next(data_iter)
+            except StopIteration:
+                if sampler: 
+                    epoch = step // len(dataloader) + 1
+                    sampler.set_epoch(epoch)
+                data_iter = iter(dataloader)
+                batch = next(data_iter)
+            
+            samples_seen += CONFIG['batch_size'] * WORLD_SIZE
+            
+            # 1. 生成 + SAS
+            experience = trainer.generate_and_prepare(
+                batch, 
+                max_gen_len=CONFIG['max_gen_len']
+            )
+            
+            step_reward = experience['rewards'].mean().item()
+            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}"
+            
+            if loss is not None:
+                if running_loss == 0: running_loss = loss
+                else: running_loss = 0.9 * running_loss + 0.1 * loss
+                status_dict["Ls"] = f"{running_loss:.3f}"
+                
+                if IS_MAIN:
+                    current_lr = trainer.optimizer.param_groups[0]['lr']
+                    metrics_data = {
+                        "step": step,
+                        "running_reward": float(running_reward),
+                        "reward": float(step_reward),
+                        "loss": float(loss),
+                        "lr": float(current_lr),
+                        "samples_seen": samples_seen,
+                        "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
+                    
+                    with open(os.path.join(CONFIG['save_dir'], "metrics.jsonl"), "a", encoding='utf-8') as f:
+                        f.write(json.dumps(metrics_data) + "\n")
+                    
+                    if step % 10 == 0:
+                        log_msg = f"Step {step} | Reward: {step_reward:.4f} | Loss: {loss:.4f}"
+                        progress_bar.write(log_msg)
+                        if file_handler:
+                            file_handler.emit(logging.LogRecord(
+                                name="train", level=logging.INFO, pathname=__file__, lineno=0,
+                                msg=log_msg, args=(), exc_info=None
+                            ))
+            else:
+                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:
+                if not is_accum_boundary:
+                    msg = "Saving checkpoint during gradient accumulation! Partial gradients will be lost."
+                    progress_bar.write(msg)
+                    if file_handler: logger.warning(msg)
+                
+                save_path = f"{CONFIG['save_dir']}/step_{step}.pt"
+                model_to_save = actor.module if hasattr(actor, 'module') else actor
+                
+                torch.save({
+                    'step': step,
+                    'samples_seen': samples_seen,
+                    'model_state_dict': model_to_save.state_dict(),
+                    'trainer_state_dict': trainer.state_dict(),
+                    'rng_state': torch.get_rng_state(),
+                    'cuda_rng_state': torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None
+                }, save_path)
+                
+                msg = f"Checkpoint saved: {save_path}"
+                progress_bar.write(msg)
+                if file_handler: logger.info(msg)
+            
+            del experience
+            del batch
+            
+        except Exception as e:
+            err_msg = f"Step {step} Error: {e}"
+            if IS_MAIN:
+                progress_bar.write(err_msg)
+            logger.error(err_msg)
+            import traceback
+            traceback.print_exc()
+            continue
+
+    if IS_MAIN:
+        final_path = f"{CONFIG['save_dir']}/final_dcpo.pt"
+        model_to_save = actor.module if hasattr(actor, 'module') else actor
+        torch.save({'model_state_dict': model_to_save.state_dict()}, final_path)
+        print("DCPO Training Finished.")
+        
+    if WORLD_SIZE > 1:
+        dist.destroy_process_group()
+
+if __name__ == "__main__":
+    main()