posttrain.py

import os
import torch
import torch.nn.functional as F
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 gc
from model import MultiModalDenseTransformer
from grpo_dataloader import create_grpo_prompt_dataloader
from data_loader import (
    create_posttrain_dataloader,
    create_preference_dataloader
)
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 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,
        model: MultiModalDenseTransformer,
        tokenizer,
        learning_rate: float = 1e-5,
        weight_decay: float = 0.01,
        num_epochs: int = 3,
        gradient_accumulation_steps: int = 1,
        max_grad_norm: float = 1.0,
        log_interval: int = 10,
        eval_interval: int = 500,
        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.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.optimizer = torch.optim.AdamW(
            self.model.parameters(),
            lr=learning_rate,
            weight_decay=weight_decay,
            betas=(0.9, 0.95),
            eps=1e-8
        )
       
        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
        self.checkpoint_dir = Path(checkpoint_dir)

        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")
        
        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, _ = 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
                position_ids[i] = positions * non_pad_mask.long()
        
        labels = input_ids.clone()
        instr_len = instruction_ids.shape[1]
        labels[:, :instr_len] = -100
        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, 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
            )
            raw_loss = loss.item()
            loss = loss / self.gradient_accumulation_steps
       
        self.scaler.scale(loss).backward()
        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
           
            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[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, 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)
                total_loss += loss.item()
                num_batches += 1
       
        self.model.train()
        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):
        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(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(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']
                    
                    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:
                        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)
                        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)
                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):
        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(),
            '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 = 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'])

        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.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)
        
        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': 2048,
        'dropout': 0.0,
        'use_moe': False,

        'batch_size': 4,
        'gradient_accumulation_steps': 16,
        'learning_rate': 1e-5,
        'weight_decay': 0.01,
        'num_epochs': 4,
        'max_grad_norm': 1.0,
        
        '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': 2048,
        'num_workers': 2,

        'do_rlhf': False,
        'preference_dataset': 'grpo_preferences_local',
        'grpo_prompt_mix': 'default',
        'grpo_iterations': 4,
        'grpo_kl_coef': 0.04,
        'grpo_group_size': 4,
        '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': 1000,
        'save_interval': 50,
    }
    
    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')

    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']  
        )
        
        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()