Update posttrain.py

posttrain.py

@@ -1,31 +1,69 @@
 import os
 import torch
 import torch.nn.functional as F
-from transformers import AutoTokenizer
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.utils.data import DataLoader, DistributedSampler
+from transformers import AutoTokenizer, get_cosine_schedule_with_warmup  # 添加 scheduler
 from pathlib import Path
 import logging
 from tqdm import tqdm
 import json
 from datetime import datetime
-import copy
+import gc
 from model import MultiModalDenseTransformer
-
+from grpo_dataloader import create_grpo_prompt_dataloader
 from data_loader import (
     create_posttrain_dataloader,
     create_preference_dataloader
 )
-from data_config import POSTTRAIN_MIX
-from reward_model import RewardModel, RewardModelTrainer
-from grpo import GRPOTrainer
+from reward_model import RewardModel 
+from grpo import GRPOZeroTrainer
 from typing import Optional
+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)
+        return rank, local_rank, world_size
+    else:
+        print("Not running in distributed mode. Fallback to single GPU.")
+        return 0, 0, 1
 
+RANK, LOCAL_RANK, WORLD_SIZE = setup_distributed()
+IS_MAIN_PROCESS = RANK == 0
 logging.basicConfig(
-    level=logging.INFO,
-    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+    level=logging.INFO if IS_MAIN_PROCESS else logging.WARNING,
+    format=f'%(asctime)s - [Rank {RANK}] - %(name)s - %(levelname)s - %(message)s'
 )
 logger = logging.getLogger(__name__)
 os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
 
+def force_cleanup():
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    if WORLD_SIZE > 1:
+        dist.barrier() 
+
+def get_distributed_dataloader(original_loader, batch_size, num_workers):
+    dataset = original_loader.dataset
+    collate_fn = original_loader.collate_fn
+   
+    sampler = DistributedSampler(dataset, shuffle=True) if WORLD_SIZE > 1 else None
+   
+    return DataLoader(
+        dataset,
+        batch_size=batch_size,
+        num_workers=num_workers,
+        pin_memory=True,
+        sampler=sampler,
+        shuffle=(sampler is None),
+        collate_fn=collate_fn
+    )
+
 class PostTrainer:
     def __init__(
         self,
@@ -38,79 +76,125 @@ class PostTrainer:
         max_grad_norm: float = 1.0,
         log_interval: int = 10,
         eval_interval: int = 500,
-        save_interval: int = 1000,
-        checkpoint_dir: str = "checkpoints/posttrain"
+        save_interval: int = 1300,
+        checkpoint_dir: str = "checkpoints/posttrain",
+        warmup_steps: int = 100,  
+        scheduler_type: str = "cosine",  
+        min_lr_ratio: float = 0.1,  
+        total_steps: Optional[int] = None  
     ):
-        self.model = model
+        self.device = torch.device(f'cuda:{LOCAL_RANK}')
+        self.model = model.to(self.device)
+
+        if WORLD_SIZE > 1:
+            self.model = DDP(self.model, device_ids=[LOCAL_RANK], output_device=LOCAL_RANK)
+       
         self.tokenizer = tokenizer
-        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-        
-        self.model.to(self.device)
-        
-        # 优化器
+       
         self.optimizer = torch.optim.AdamW(
-            model.parameters(),
+            self.model.parameters(),
             lr=learning_rate,
             weight_decay=weight_decay,
             betas=(0.9, 0.95),
             eps=1e-8
         )
-        
-        # 混合精度
-        self.use_amp = torch.cuda.is_available()
+       
+        self.use_amp = True
         self.scaler = torch.amp.GradScaler('cuda', enabled=self.use_amp)
-        
-        # 训练参数
+       
         self.num_epochs = num_epochs
         self.gradient_accumulation_steps = gradient_accumulation_steps
         self.max_grad_norm = max_grad_norm
         self.log_interval = log_interval
         self.eval_interval = eval_interval
         self.save_interval = save_interval
-        
-        # Checkpoint管理
         self.checkpoint_dir = Path(checkpoint_dir)
-        self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
-        
-        # 训练状态
+
+        self.warmup_steps = warmup_steps
+        self.scheduler_type = scheduler_type
+        self.min_lr_ratio = min_lr_ratio
+        self.learning_rate = learning_rate
+        self.total_steps = total_steps
+        self.scheduler = None  
+       
+        if IS_MAIN_PROCESS:
+            self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
+            log_file_name = f"train_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log"
+            log_path = self.checkpoint_dir / log_file_name
+            file_handler = logging.FileHandler(log_path, encoding='utf-8')
+            file_handler.setLevel(logging.INFO)
+            formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
+            file_handler.setFormatter(formatter)
+            logger.addHandler(file_handler)
+            
         self.global_step = 0
         self.best_eval_loss = float('inf')
+
+    def _create_scheduler(self, total_steps: int):
+        if self.scheduler_type == "cosine":
+            from transformers import get_cosine_schedule_with_warmup
+            scheduler = get_cosine_schedule_with_warmup(
+                self.optimizer,
+                num_warmup_steps=self.warmup_steps,
+                num_training_steps=total_steps,
+                num_cycles=0.5  
+            )
+        elif self.scheduler_type == "linear":
+            from transformers import get_linear_schedule_with_warmup
+            scheduler = get_linear_schedule_with_warmup(
+                self.optimizer,
+                num_warmup_steps=self.warmup_steps,
+                num_training_steps=total_steps
+            )
+        elif self.scheduler_type == "constant":
+            from transformers import get_constant_schedule_with_warmup
+            scheduler = get_constant_schedule_with_warmup(
+                self.optimizer,
+                num_warmup_steps=self.warmup_steps
+            )
+        elif self.scheduler_type == "cosine_with_min_lr":
+            from transformers import get_cosine_schedule_with_warmup
+            scheduler = get_cosine_schedule_with_warmup(
+                self.optimizer,
+                num_warmup_steps=self.warmup_steps,
+                num_training_steps=total_steps,
+                num_cycles=0.5
+            )
+            scheduler = MinLRSchedulerWrapper(
+                scheduler, 
+                self.optimizer, 
+                min_lr=self.learning_rate * self.min_lr_ratio
+            )
+        else:
+            raise ValueError(f"Unknown scheduler type: {self.scheduler_type}")
+        
+        if IS_MAIN_PROCESS:
+            logger.info(f"Created {self.scheduler_type} scheduler with {self.warmup_steps} warmup steps and {total_steps} total steps")
         
-        logger.info(f"PostTrainer initialized:")
-        logger.info(f"  Device: {self.device}")
-        logger.info(f"  Learning Rate: {learning_rate}")
-        logger.info(f"  Num Epochs: {num_epochs}")
-        logger.info(f"  Gradient Accumulation: {gradient_accumulation_steps}")
+        return scheduler
 
     def train_step(self, batch: dict) -> dict:
-        """单步训练"""
         instruction_ids = batch['instruction'].to(self.device)
         response_ids = batch['response'].to(self.device)
-        
         instruction_mask = batch['instruction_mask'].to(self.device)
         response_mask = batch['response_mask'].to(self.device)
-        
+       
         input_ids = torch.cat([instruction_ids, response_ids], dim=1)
         attention_mask = torch.cat([instruction_mask, response_mask], dim=1)
-        
-        batch_size , seq_len = input_ids.shape
-        position_ids=torch.zeros_like(input_ids)
-
+       
+        batch_size, _ = input_ids.shape
+        position_ids = torch.zeros_like(input_ids)
         for i in range(batch_size):
             non_pad_mask = attention_mask[i].bool()
             if non_pad_mask.any():
-                positions=torch.cumsum(non_pad_mask.long(), dim=0) -1 
+                positions = torch.cumsum(non_pad_mask.long(), dim=0) - 1
                 position_ids[i] = positions * non_pad_mask.long()
-        labels = input_ids.clone()
         
-        # 屏蔽 Instruction 部分
+        labels = input_ids.clone()
         instr_len = instruction_ids.shape[1]
         labels[:, :instr_len] = -100
-
         labels[attention_mask == 0] = -100
- 
         
-        # 准备输入数据
         input_data = {
             'segments': [{
                 'type': 'text',
@@ -118,18 +202,14 @@ class PostTrainer:
                 'modality_id': 0
             }]
         }
-        
-        # 前向传播
+       
         with torch.amp.autocast('cuda', enabled=self.use_amp):
-            outputs = self.model(input_data, attention_mask=attention_mask,
-                                 position_ids = position_ids)
-            
+            outputs = self.model(input_data, attention_mask=attention_mask, position_ids=position_ids)
             logits = outputs['logits']
-            
-            # 计算损失
+           
             shift_logits = logits[:, :-1, :].contiguous()
             shift_labels = labels[:, 1:].contiguous()
-            
+           
             loss = F.cross_entropy(
                 shift_logits.view(-1, shift_logits.size(-1)),
                 shift_labels.view(-1),
@@ -137,399 +217,587 @@ class PostTrainer:
             )
             raw_loss = loss.item()
             loss = loss / self.gradient_accumulation_steps
-        
-        # 反向传播
+       
         self.scaler.scale(loss).backward()
-        
-        return {
-            'loss': raw_loss
-        }
+        return {'loss': raw_loss}
 
     def optimizer_step(self):
-        """优化器步骤"""
         self.scaler.unscale_(self.optimizer)
         grad_norm = torch.nn.utils.clip_grad_norm_(
             self.model.parameters(),
             self.max_grad_norm
         )
-        
         self.scaler.step(self.optimizer)
         self.scaler.update()
+        if self.scheduler is not None:
+            self.scheduler.step()
+        
         self.optimizer.zero_grad(set_to_none=True)
         self.global_step += 1
         return grad_norm.item()
 
     @torch.no_grad()
     def evaluate(self, dataloader, max_batches: int = 50) -> float:
-        """评估"""
         self.model.eval()
         total_loss = 0.0
         num_batches = 0
-        
+       
         for i, batch in enumerate(dataloader):
-            if i >= max_batches:
-                break
-            
-            if batch is None:
-                continue
-            
+            if i >= max_batches: break
+            if batch is None: continue
+           
             instruction_ids = batch['instruction'].to(self.device)
             response_ids = batch['response'].to(self.device)
             input_ids = torch.cat([instruction_ids, response_ids], dim=1)
+
+            instruction_mask = batch['instruction_mask'].to(self.device)
+            response_mask = batch['response_mask'].to(self.device)
+            attention_mask = torch.cat([instruction_mask, response_mask], dim=1)
+
+            position_ids = torch.zeros_like(input_ids)
+            for i in range(input_ids.shape[0]):
+                non_pad = attention_mask[i].bool()
+                if non_pad.any():
+                    position_ids[i] = (torch.cumsum(non_pad.long(), dim=0) - 1) * non_pad.long()
             
             labels = input_ids.clone()
             labels[:, :instruction_ids.shape[1]] = -100
-            labels[input_ids == self.tokenizer.pad_token_id] = -100
-            
-            input_data = {
-                'segments': [{
-                    'type': 'text',
-                    'data': input_ids,
-                    'modality_id': 0
-                }]
-            }
-            
+            labels[attention_mask == 0] = -100
+            input_data = {'segments': [{'type': 'text', 'data': input_ids, 'modality_id': 0}]}
+
             with torch.amp.autocast('cuda', enabled=self.use_amp):
-                outputs = self.model(input_data)
+                outputs = self.model(input_data, attention_mask=attention_mask, position_ids=position_ids)
                 logits = outputs['logits']
-                
                 shift_logits = logits[:, :-1, :].contiguous()
                 shift_labels = labels[:, 1:].contiguous()
-                
-                loss = F.cross_entropy(
-                    shift_logits.view(-1, shift_logits.size(-1)),
-                    shift_labels.view(-1),
-                    ignore_index=-100
-                )
-                
+                loss = F.cross_entropy(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1), ignore_index=-100)
                 total_loss += loss.item()
                 num_batches += 1
-        
+       
         self.model.train()
-        return total_loss / max(num_batches, 1)
+        avg_loss = total_loss / max(num_batches, 1)
+       
+        if WORLD_SIZE > 1:
+            loss_tensor = torch.tensor(avg_loss).to(self.device)
+            dist.all_reduce(loss_tensor, op=dist.ReduceOp.AVG)
+            avg_loss = loss_tensor.item()
+           
+        return avg_loss
 
-    def train(
-        self,
-        train_dataloader,
-        eval_dataloader=None,
-        resume_from: Optional[str] = None
-    ):
-        """训练循环"""
-        logger.info("\n" + "="*80)
-        logger.info("Starting Post-Training (SFT)")
-        logger.info("="*80 + "\n")
-        
+    def train(self, train_dataloader, eval_dataloader=None, resume_from: Optional[str] = None):
+        if IS_MAIN_PROCESS:
+            logger.info("Starting Post-Training (SFT) with LR Scheduler - DDP Mode")
+       
+        if self.total_steps is None:
+            steps_per_epoch = len(train_dataloader) // self.gradient_accumulation_steps
+            self.total_steps = steps_per_epoch * self.num_epochs
+            if IS_MAIN_PROCESS:
+                logger.info(f"Calculated total training steps: {self.total_steps}")
+        
+        self.scheduler = self._create_scheduler(self.total_steps)
+        
+        start_epoch = 0
         if resume_from:
             self.load_checkpoint(resume_from)
+            steps_per_epoch = len(train_dataloader) // self.gradient_accumulation_steps
+            start_epoch = self.global_step // steps_per_epoch
+            if IS_MAIN_PROCESS:
+                logger.info(f"Resuming training from epoch {start_epoch}, global step {self.global_step}")
         
         self.model.train()
-        
-        for epoch in range(self.num_epochs):
-            logger.info(f"\nEpoch {epoch+1}/{self.num_epochs}")
-            
-            progress_bar = tqdm(train_dataloader, desc=f"Epoch {epoch+1}")
+       
+        for epoch in range(start_epoch, self.num_epochs):
+            if hasattr(train_dataloader.sampler, 'set_epoch'):
+                train_dataloader.sampler.set_epoch(epoch)
+               
+            if IS_MAIN_PROCESS:
+                logger.info(f"\nEpoch {epoch+1}/{self.num_epochs}")
+
+            iterator = tqdm(train_dataloader, desc=f"Epoch {epoch+1}", disable=not IS_MAIN_PROCESS)
+           
             running_loss = 0.0
             step_in_accumulation = 0
-            
-            for batch_idx, batch in enumerate(progress_bar):
-                if batch is None:
+           
+            for batch_idx, batch in enumerate(iterator):
+                if batch is None: continue
+
+                if 'instruction' not in batch or 'response' not in batch:
+                    if IS_MAIN_PROCESS:
+                        logger.warning(f"Skipping invalid batch at index {batch_idx}")
                     continue
                 
-                # 训练步骤
                 stats = self.train_step(batch)
                 running_loss += stats['loss']
                 step_in_accumulation += 1
-                
-                # 优化器更新
+               
                 if step_in_accumulation == self.gradient_accumulation_steps:
                     grad_norm = self.optimizer_step()
                     step_in_accumulation = 0
+
+                    current_lr = self.optimizer.param_groups[0]['lr']
                     
-                    # 更新进度条
-                    progress_bar.set_postfix({'loss': f"{stats['loss']:.4f}"})
-                    
-                    # 日志
+                    if IS_MAIN_PROCESS:
+                        iterator.set_postfix({
+                            'loss': f"{stats['loss']:.4f}",
+                            'lr': f"{current_lr:.2e}"
+                        })
+                   
                     if self.global_step % self.log_interval == 0:
-                        avg_loss = running_loss / self.log_interval
-                        logger.info(
-                            f"Step {self.global_step} | "
-                            f"Epoch {epoch+1} | "
-                            f"Loss: {avg_loss:.4f}"
-                        )
+                        current_loss_tensor = torch.tensor(running_loss).to(self.device)
+                        if WORLD_SIZE > 1:
+                            dist.all_reduce(current_loss_tensor, op=dist.ReduceOp.AVG)
+                        avg_loss = current_loss_tensor.item() / (self.log_interval * self.gradient_accumulation_steps)
+                       
+                        if IS_MAIN_PROCESS:
+                            logger.info(
+                                f"Step: {self.global_step} | "
+                                f"Loss: {avg_loss:.6f} | "
+                                f"GradNorm: {grad_norm:.4f} | "
+                                f"LR: {current_lr:.2e} | "
+                                f"Progress: {self.global_step}/{self.total_steps}"
+                            )
                         running_loss = 0.0
-                    
-                    # 评估
+                   
                     if eval_dataloader and self.global_step % self.eval_interval == 0:
                         eval_loss = self.evaluate(eval_dataloader)
-                        logger.info(f"Eval Loss: {eval_loss:.4f}")
-                        
-                        if eval_loss < self.best_eval_loss:
-                            self.best_eval_loss = eval_loss
-                            self.save_checkpoint(
-                                self.checkpoint_dir / "best_model.pt",
-                                is_best=True
-                            )
-                    
-                    # 保存
-                    if self.global_step % self.save_interval == 0:
-                        self.save_checkpoint(
-                            self.checkpoint_dir / f"step_{self.global_step}.pt"
-                        )
-            
-            # Epoch结束评估
+                        if IS_MAIN_PROCESS:
+                            logger.info(f"Eval Loss: {eval_loss:.4f}")
+                            if eval_loss < self.best_eval_loss:
+                                self.best_eval_loss = eval_loss
+                                self.save_checkpoint(self.checkpoint_dir / "best_model.pt", is_best=True)
+                   
+                    if self.global_step % self.save_interval == 0 and IS_MAIN_PROCESS:
+                        self.save_checkpoint(self.checkpoint_dir / f"step_{self.global_step}.pt")
+           
             if eval_dataloader:
                 eval_loss = self.evaluate(eval_dataloader)
-                logger.info(f"\nEpoch {epoch+1} Eval Loss: {eval_loss:.4f}")
-        
-        logger.info("\n" + "="*80)
-        logger.info("Post-Training Complete!")
-        logger.info(f"  Best Eval Loss: {self.best_eval_loss:.4f}")
-        logger.info("="*80 + "\n")
-        
-        self.save_checkpoint(self.checkpoint_dir / "final_model.pt")
+                if IS_MAIN_PROCESS:
+                    logger.info(f"\nEpoch {epoch+1} Eval Loss: {eval_loss:.4f}")
+       
+        if IS_MAIN_PROCESS:
+            self.save_checkpoint(self.checkpoint_dir / "final_model.pt")
 
     def save_checkpoint(self, path: Path, is_best: bool = False):
-        """保存checkpoint"""
+        if not IS_MAIN_PROCESS: return
+        model_to_save = self.model.module if hasattr(self.model, 'module') else self.model
         checkpoint = {
-            'model_state_dict': self.model.state_dict(),
+            'model_state_dict': model_to_save.state_dict(),
             'optimizer_state_dict': self.optimizer.state_dict(),
+            'scheduler_state_dict': self.scheduler.state_dict() if self.scheduler else None,  # 新增
             'scaler_state_dict': self.scaler.state_dict() if self.use_amp else None,
             'global_step': self.global_step,
             'best_eval_loss': self.best_eval_loss,
             'timestamp': datetime.now().isoformat()
         }
-        
         torch.save(checkpoint, path)
         logger.info(f"Checkpoint saved to {path}" + (" (BEST)" if is_best else ""))
 
     def load_checkpoint(self, path: str):
-        """加载checkpoint"""
         checkpoint = torch.load(path, map_location=self.device)
-        
-        self.model.load_state_dict(checkpoint['model_state_dict'])
+       
+        model_to_load = self.model.module if hasattr(self.model, 'module') else self.model
+        model_to_load.load_state_dict(checkpoint['model_state_dict'])
         self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+
+        if self.scheduler and checkpoint.get('scheduler_state_dict'):
+            self.scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
         
         if self.use_amp and checkpoint.get('scaler_state_dict'):
             self.scaler.load_state_dict(checkpoint['scaler_state_dict'])
-        
         self.global_step = checkpoint['global_step']
-        self.best_eval_loss = checkpoint['best_eval_loss']
+        self.best_eval_loss = checkpoint.get('best_eval_loss', float('inf'))
+        if IS_MAIN_PROCESS:
+            logger.info(f"Checkpoint loaded from {path}")
+
+
+class MinLRSchedulerWrapper:
+    def __init__(self, scheduler, optimizer, min_lr):
+        self.scheduler = scheduler
+        self.optimizer = optimizer
+        self.min_lr = min_lr
+    
+    def step(self):
+        self.scheduler.step()
+        for param_group in self.optimizer.param_groups:
+            param_group['lr'] = max(param_group['lr'], self.min_lr)
+    
+    def state_dict(self):
+        return {
+            'scheduler': self.scheduler.state_dict(),
+            'min_lr': self.min_lr
+        }
+    
+    def load_state_dict(self, state_dict):
+        self.scheduler.load_state_dict(state_dict['scheduler'])
+        self.min_lr = state_dict['min_lr']
+
+class RewardTrainer:
+    def __init__(
+        self,
+        tokenizer,
+        reward_model: RewardModel,
+        learning_rate: float = 1e-5,
+        weight_decay: float = 0.01,
+        num_epochs: int = 1,
+        gradient_accumulation_steps: int = 8,
+        max_grad_norm: float = 1.0,
+        log_interval: int = 10,
+        save_interval: int = 2000,
+        checkpoint_dir: str = "checkpoints/reward_checkpoints"
+    ):
+        self.device = torch.device(f'cuda:{LOCAL_RANK}')
+        self.model = reward_model  
+        self.tokenizer = tokenizer
+        self.pad_token_id = tokenizer.pad_token_id
+        self.optimizer = torch.optim.AdamW(
+            self.model.parameters(),
+            lr=learning_rate,
+            weight_decay=weight_decay
+        )
+        self.use_amp = True
+        self.scaler = torch.amp.GradScaler('cuda', enabled=self.use_amp)
+        self.num_epochs = num_epochs
+        self.gradient_accumulation_steps = gradient_accumulation_steps
+        self.max_grad_norm = max_grad_norm
+        self.log_interval = log_interval
+        self.save_interval = save_interval
+        self.checkpoint_dir = Path(checkpoint_dir)
         
-        logger.info(f"Checkpoint loaded from {path}")
+        self.file_handler = None
+        if IS_MAIN_PROCESS:
+            self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
+            log_file_name = f"reward_train_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log"
+            log_path = self.checkpoint_dir / log_file_name
+            
+            self.file_handler = logging.FileHandler(log_path, encoding='utf-8')
+            self.file_handler.setLevel(logging.INFO)
+            formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
+            self.file_handler.setFormatter(formatter)
+            logger.addHandler(self.file_handler)
+
+        self.global_step = 0
+        self.running_loss = 0.0
+        self.running_acc = 0.0
+
+    def train_step(self, batch: dict) -> dict:
+        chosen_ids = batch['chosen'].to(self.device)
+        rejected_ids = batch['rejected'].to(self.device)
+
+        batch_size = chosen_ids.size(0)
+
+        chosen_attention_mask = (chosen_ids != self.pad_token_id).long()
+        chosen_position_ids = torch.cumsum(chosen_attention_mask, dim=1) - 1
+        chosen_position_ids = chosen_position_ids * chosen_attention_mask
+
+        chosen_input = {'segments': [{'type': 'text', 'data': chosen_ids, 'modality_id': 0}]}
+
+        rejected_attention_mask = (rejected_ids != self.pad_token_id).long()
+        rejected_position_ids = torch.cumsum(rejected_attention_mask, dim=1) - 1
+        rejected_position_ids = rejected_position_ids * rejected_attention_mask
+
+        rejected_input = {'segments': [{'type': 'text', 'data': rejected_ids, 'modality_id': 0}]}
+
+        with torch.amp.autocast('cuda', enabled=self.use_amp):
+            chosen_rewards_full = self.model(
+                chosen_input,
+                attention_mask=chosen_attention_mask,
+                position_ids=chosen_position_ids
+            )
+
+            rejected_rewards_full = self.model(
+                rejected_input,
+                attention_mask=rejected_attention_mask,
+                position_ids=rejected_position_ids
+            )
+
+            chosen_last_idx = chosen_attention_mask.sum(dim=1) - 1
+            rejected_last_idx = rejected_attention_mask.sum(dim=1) - 1
+
+            chosen_rewards = chosen_rewards_full[torch.arange(batch_size, device=self.device), chosen_last_idx]
+            rejected_rewards = rejected_rewards_full[torch.arange(batch_size, device=self.device), rejected_last_idx]
+
+            loss = -F.logsigmoid(chosen_rewards - rejected_rewards).mean()
+            acc = (chosen_rewards > rejected_rewards).float().mean().item()
+
+            loss = loss / self.gradient_accumulation_steps
+
+        self.scaler.scale(loss).backward()
+        raw_loss = loss.item() * self.gradient_accumulation_steps
+        return {'loss': raw_loss, 'acc': acc}
+
+    def optimizer_step(self):
+        self.scaler.unscale_(self.optimizer)
+        grad_norm = torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.max_grad_norm)
+        self.scaler.step(self.optimizer)
+        self.scaler.update()
+        self.optimizer.zero_grad(set_to_none=True)
+        self.global_step += 1
+        return grad_norm.item()
+
+    def save_checkpoint(self, path: Path):
+        if not IS_MAIN_PROCESS:
+            return
+        model_to_save = self.model.module if hasattr(self.model, 'module') else self.model
+        checkpoint = {
+            'model_state_dict': model_to_save.state_dict(),
+            'optimizer_state_dict': self.optimizer.state_dict(),
+            'scaler_state_dict': self.scaler.state_dict(),
+            'global_step': self.global_step,
+        }
+        torch.save(checkpoint, path)
+        logger.info(f"Reward checkpoint saved: {path}")
+
+    def load_checkpoint(self, path: str):
+        checkpoint = torch.load(path, map_location=self.device)
+        model_to_load = self.model.module if hasattr(self.model, 'module') else self.model
+        model_to_load.load_state_dict(checkpoint['model_state_dict'])
+        self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+        self.scaler.load_state_dict(checkpoint['scaler_state_dict'])
+        self.global_step = checkpoint['global_step']
+        if IS_MAIN_PROCESS:
+            logger.info(f"Reward checkpoint loaded: {path} (step {self.global_step})")
+
+    def train(self, dataloader, resume_from: Optional[str] = None):
+        try:
+            if resume_from:
+                self.load_checkpoint(resume_from)
+
+            self.model.train()
+            for epoch in range(self.num_epochs):
+                if hasattr(dataloader.sampler, 'set_epoch'):
+                    dataloader.sampler.set_epoch(epoch)
+
+                iterator = tqdm(dataloader, desc=f"Reward Epoch {epoch+1}/{self.num_epochs}", disable=not IS_MAIN_PROCESS)
+                accum_steps = 0
+                self.running_loss = 0.0
+                self.running_acc = 0.0
+
+                for batch in iterator:
+                    if batch is None or 'chosen' not in batch:
+                        continue
+
+                    stats = self.train_step(batch)
+                    single_step_loss = stats['loss'] / self.gradient_accumulation_steps
+                    self.running_loss += single_step_loss
+                    self.running_acc += stats['acc']
+                    accum_steps += 1
+
+                    if accum_steps == self.gradient_accumulation_steps:
+                        grad_norm = self.optimizer_step()
+                        accum_steps = 0
+
+                        if IS_MAIN_PROCESS:
+                            iterator.set_postfix({'loss': f"{stats['loss']:.4f}", 'acc': f"{stats['acc']:.4f}"})
+
+                        if self.global_step % self.log_interval == 0:
+                            avg_loss = self.running_loss / self.log_interval
+                            avg_acc = self.running_acc / self.log_interval
+                            if WORLD_SIZE > 1:
+                                loss_tensor = torch.tensor(avg_loss, device=self.device)
+                                acc_tensor = torch.tensor(avg_acc, device=self.device)
+                                dist.all_reduce(loss_tensor, op=dist.ReduceOp.AVG)
+                                dist.all_reduce(acc_tensor, op=dist.ReduceOp.AVG)
+                                avg_loss = loss_tensor.item()
+                                avg_acc = acc_tensor.item()
+                            if IS_MAIN_PROCESS:
+                                logger.info(f"Reward Step {self.global_step} | Loss {avg_loss:.6f} | Acc {avg_acc:.4f} | Grad {grad_norm:.4f}")
+                            self.running_loss = 0.0
+                            self.running_acc = 0.0
+
+                        if self.global_step % self.save_interval == 0 and self.global_step > 0:
+                            self.save_checkpoint(self.checkpoint_dir / f"step_{self.global_step}.pt")
+        finally:
+            if IS_MAIN_PROCESS and self.file_handler:
+                logger.removeHandler(self.file_handler)
+                self.file_handler.close()
+                self.file_handler = None
+
+
+def load_checkpoint_flexible(model, path, device):
+    logger.info(f"Loading weights from {path}...")
+    checkpoint = torch.load(path, map_location=device)
+    
+    state_dict = None
+    if 'actor_state_dict' in checkpoint:
+        logger.info("Detected GRPO checkpoint format.")
+        state_dict = checkpoint['actor_state_dict']
+
+    elif 'model_state_dict' in checkpoint:
+        logger.info("Detected Standard/SFT checkpoint format.")
+        state_dict = checkpoint['model_state_dict']
+
+    else:
+        logger.info("Detected raw state dict format.")
+        state_dict = checkpoint
+
+    model_has_module = hasattr(model, 'module')
+    
+    new_state_dict = {}
+    for k, v in state_dict.items():
+        if k.startswith('module.') and not model_has_module:
+            new_state_dict[k[7:]] = v
+        else:
+            new_state_dict[k] = v
+            
+    model.load_state_dict(new_state_dict, strict=False)
+    logger.info("Weights loaded successfully.")
+    del checkpoint
+    gc.collect()
+    torch.cuda.empty_cache()
 
 def main():
-    """主函数"""
-    # 配置
     config = {
-        # 模型配置
         'model_dim': 1536,
         'vocab_size': 151665,
         'n_layers': 12,
         'n_heads': 12,
         'n_kv_heads': 4,
-        'max_seq_len': 512,
+        'max_seq_len': 2048,
         'dropout': 0.0,
         'use_moe': False,
-        # 训练配置
-        'batch_size': 2,
-        'gradient_accumulation_steps': 8,
+
+        'batch_size': 4,
+        'gradient_accumulation_steps': 16,
         'learning_rate': 1e-5,
         'weight_decay': 0.01,
-        'num_epochs': 3,
+        'num_epochs': 4,
         'max_grad_norm': 1.0,
         
-        # 数据配置
-        'data_mix': 'simple_instruct',
-        'max_samples_train': 20000,
+        'warmup_steps': 100,  
+        'scheduler_type': 'cosine', 
+        'min_lr_ratio': 0.1,  
+
+        'data_mix': 'think_math_mix',
+        'max_samples_train': None,
         'max_samples_eval': 1000,
-        'max_length': 512,
-        'num_workers': 4,
-        
-        # RLHF配置
+        'max_length': 2048,
+        'num_workers': 2,
+
         'do_rlhf': False,
-        'preference_dataset': 'hh_rlhf',
-        'grpo_iterations': 3,
+        'preference_dataset': 'grpo_preferences_local',
+        'grpo_prompt_mix': 'default',
+        'grpo_iterations': 4,
         'grpo_kl_coef': 0.04,
         'grpo_group_size': 4,
-        
-        # 路径
-        'pretrain_checkpoint': '/root/multimodal/checkpoints/pretrain_fixed/step_10000.pt',
-        'checkpoint_dir': 'checkpoints/posttrain',
+        'grpo_max_gen_len': 256,
+        'grpo_temperature': 0.9,
+        'grpo_prompt_batch_size': 1,
+        'grpo_max_prompts': 5000,
+        'grpo_resume_path': None,
+
+        'pretrain_checkpoint': '/root/checkpoints/pretrain_fixed/step_45000.pt',
+        'sft_checkpoint': '/root/checkpoints/dcpo_posttrain_round3/step_7800.pt',
+        'checkpoint_dir': '/root/checkpoints/dcpo_posttrain_round3',
         'log_interval': 50,
-        'eval_interval': 500,
-        'save_interval': 1000,
+        'eval_interval': 1000,
+        'save_interval': 50,
     }
-
-    logger.info("Configuration:")
-    logger.info(json.dumps(config, indent=2))
-
-    # 初始化tokenizer
+    
+    if IS_MAIN_PROCESS:
+        logger.info("Configuration:")
+        logger.info(json.dumps(config, indent=2))
+        logger.info(f"Running DDP on Rank: {RANK}, Local Rank: {LOCAL_RANK}, World Size: {WORLD_SIZE}")
+    
     logger.info("\nInitializing tokenizer...")
     tokenizer = AutoTokenizer.from_pretrained(
         "Qwen/Qwen2.5-7B-Instruct",
         use_fast=True,
         trust_remote_code=True
     )
-
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
         tokenizer.pad_token_id = tokenizer.eos_token_id
-
+    
     config['vocab_size'] = len(tokenizer)
+    
+    def create_model_architecture():
+        return MultiModalDenseTransformer(
+            model_dim=config['model_dim'],
+            vocab_size=config['vocab_size'],
+            n_layers=config['n_layers'],
+            n_heads=config['n_heads'],
+            n_kv_heads=config['n_kv_heads'],
+            max_seq_len=config['max_seq_len'],
+            dropout=config['dropout'],
+            use_moe=config['use_moe'],
+            use_gradient_checkpointing=True,
+            rope_scaling_type="yarn",
+            use_multimodal_fusion=False,
+            use_contrastive=False
+        )
+    
+    checkpoint_to_load = config.get('sft_checkpoint') or config.get('pretrain_checkpoint')
 
-    # 初始化或加载模型
-    logger.info("\nInitializing model...")
-    model = MultiModalDenseTransformer(
-        model_dim=config['model_dim'],
-        vocab_size=config['vocab_size'],
-        n_layers=config['n_layers'],
-        n_heads=config['n_heads'],
-        n_kv_heads=config['n_kv_heads'],
-        max_seq_len=config['max_seq_len'],
-        dropout=config['dropout'],
-        use_moe=config['use_moe'],
-        use_gradient_checkpointing=False,
-        rope_scaling_type="yarn",
-        use_multimodal_fusion=False,
-        use_contrastive=False
-    )
-
-    if config['pretrain_checkpoint']:
-        logger.info(f"Loading pretrain checkpoint: {config['pretrain_checkpoint']}")
-        checkpoint = torch.load(config['pretrain_checkpoint'])
-        model.load_state_dict(checkpoint['model_state_dict'])
-
-    logger.info("\n" + "="*80)
-    logger.info("PHASE 1: Supervised Fine-Tuning")
-    logger.info("="*80)
-
-    # 创建数据加载器
-    train_dataloader = create_posttrain_dataloader(
-        mix_name=config['data_mix'],
-        tokenizer=tokenizer,
-        batch_size=config['batch_size'],
-        num_workers=config['num_workers'],
-        max_length=config['max_length'],
-        max_samples=config['max_samples_train'],
-        split='train',
-        shuffle=True
-    )
-
-    eval_dataloader = create_posttrain_dataloader(
-        mix_name=config['data_mix'],
-        tokenizer=tokenizer,
-        batch_size=config['batch_size'] * 2,
-        num_workers=config['num_workers'],
-        max_length=config['max_length'],
-        max_samples=config['max_samples_eval'],
-        split='train',  # 使用train的后部分作为验证
-        shuffle=False
-    )
-
-    # 创建训练器
-    trainer = PostTrainer(
-        model=model,
-        tokenizer=tokenizer,
-        learning_rate=config['learning_rate'],
-        weight_decay=config['weight_decay'],
-        num_epochs=config['num_epochs'],
-        gradient_accumulation_steps=config['gradient_accumulation_steps'],
-        max_grad_norm=config['max_grad_norm'],
-        log_interval=config['log_interval'],
-        eval_interval=config['eval_interval'],
-        save_interval=config['save_interval'],
-        checkpoint_dir=config['checkpoint_dir']
-    )
-
-    trainer.train(train_dataloader, eval_dataloader)
-
-    if config['do_rlhf']:
-        logger.info("\n" + "="*80)
-        logger.info("PHASE 2: RLHF with GRPO")
-        logger.info("="*80)
+    do_sft =   config.get('sft_checkpoint')
+   
+    if do_sft:
+        if IS_MAIN_PROCESS:
+            logger.info("\n" + "="*80)
+            logger.info("PHASE 1: Supervised Fine-Tuning with LR Scheduler")
+            logger.info("="*80)
+       
+        model = create_model_architecture()
+        if checkpoint_to_load:
+            if IS_MAIN_PROCESS: logger.info(f"Loading checkpoint for SFT: {checkpoint_to_load}")
+            checkpoint = torch.load(checkpoint_to_load, map_location=f'cuda:{LOCAL_RANK}')
+            model.load_state_dict(checkpoint['model_state_dict'])
+            del checkpoint
+        
+        _tmp_loader = create_posttrain_dataloader(
+            mix_name=config['data_mix'], tokenizer=tokenizer,
+            batch_size=config['batch_size'], num_workers=config['num_workers'],
+            max_length=config['max_length'], max_samples=config['max_samples_train'],
+            split='train', shuffle=True
+        )
+        train_dataloader = get_distributed_dataloader(_tmp_loader, config['batch_size'], config['num_workers'])
+       
+        trainer = PostTrainer(
+            model=model, 
+            tokenizer=tokenizer,
+            learning_rate=config['learning_rate'], 
+            weight_decay=config['weight_decay'],
+            num_epochs=config['num_epochs'], 
+            gradient_accumulation_steps=config['gradient_accumulation_steps'],
+            max_grad_norm=config['max_grad_norm'],
+            checkpoint_dir=config['checkpoint_dir'],
+            warmup_steps=config['warmup_steps'],  
+            scheduler_type=config['scheduler_type'],  
+            min_lr_ratio=config['min_lr_ratio']  
+        )
         
-        try:
-            # 训练奖励模型
-            logger.info("\nTraining Reward Model...")
-            
-            reward_base_model = copy.deepcopy(model)
-            reward_model = RewardModel(reward_base_model, use_value_head=True)
-            
-            preference_dataloader = create_preference_dataloader(
-                dataset_name=config['preference_dataset'],
-                tokenizer=tokenizer,
-                batch_size=config['batch_size'],
-                num_workers=config['num_workers'],
-                max_samples=5000,
-                split='train'
-            )
-            
-            reward_trainer = RewardModelTrainer(
-                reward_model=reward_model,
-                learning_rate=1e-5
-            )
-            
-            reward_trainer.train(preference_dataloader, num_epochs=1)
-            
-            # GRPO训练
-            logger.info("\nStarting GRPO Training...")
-            
-            ref_model = copy.deepcopy(model)
-            ref_model.eval()
-            
-            grpo_trainer = GRPOTrainer(
-                actor_model=model,
-                reward_model=reward_model,
-                ref_model=ref_model,
-                tokenizer=tokenizer,
-                learning_rate=1e-6,
-                kl_coef=config['grpo_kl_coef'],
-                group_size=config['grpo_group_size'],
-                update_batch_size=2,
-                use_amp=True
-            )
-            
-            # 准备prompts
-            prompt_dataloader = create_posttrain_dataloader(
-                mix_name=config['data_mix'],
-                tokenizer=tokenizer,
-                batch_size=4,
-                num_workers=2,
-                max_samples=1000,
-                split='train'
-            )
-            
-            # 提取prompts
-            prompts = []
-            for batch in prompt_dataloader:
-                if batch and batch.get('instruction') is not None:
-                    prompts.append(batch['instruction'])
-                if len(prompts) >= 200:
-                    break
-            
-            if prompts:
-                prompt_tensor = torch.cat(prompts[:200], dim=0)
-                from torch.utils.data import TensorDataset, DataLoader
-                prompt_loader = DataLoader(
-                    TensorDataset(prompt_tensor),
-                    batch_size=4
-                )
-                
-                grpo_trainer.train(
-                    prompt_loader,
-                    num_iterations=config['grpo_iterations'],
-                    max_gen_len=50,
-                    save_path=config['checkpoint_dir'] + "/grpo"
-                )
-            
-        except Exception as e:
-            logger.error(f"Error in RLHF: {e}")
-            import traceback
-            traceback.print_exc()
-
-    logger.info("\n" + "="*80)
-    logger.info("All Training Complete!")
-    logger.info("="*80)
+        sft_resume_path = None
+        if IS_MAIN_PROCESS:
+            checkpoint_dir = Path(config['checkpoint_dir'])
+            if checkpoint_dir.exists():
+                ckpts = sorted([p for p in checkpoint_dir.glob("step_*.pt")], key=lambda p: int(p.stem.split('_')[1]))
+                if ckpts:
+                    latest = ckpts[-1]
+                    sft_resume_path = str(latest)
+                    logger.info(f"Resuming SFT training from {sft_resume_path}")
+        
+        if WORLD_SIZE > 1:
+            if IS_MAIN_PROCESS:
+                resume_path_list = [sft_resume_path]
+            else:
+                resume_path_list = [None]
+            dist.broadcast_object_list(resume_path_list, src=0)
+            sft_resume_path = resume_path_list[0]
+        
+        trainer.train(train_dataloader, None, resume_from=sft_resume_path)
+       
+        sft_save_path = Path(config['checkpoint_dir']) / "sft_complete.pt"
+        trainer.save_checkpoint(sft_save_path)
+       
+        checkpoint_to_load = str(sft_save_path)
+       
+        del model, trainer, train_dataloader
+        force_cleanup()
+    
+    if IS_MAIN_PROCESS:
+        logger.info("Training Complete!")
+   
+    if WORLD_SIZE > 1:
+        dist.destroy_process_group()
 
 if __name__ == "__main__":
     main()