grpo_r1_train.py

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
from model import MultiModalDenseTransformer
from grpo import GRPOZeroTrainer

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)
        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

logging.basicConfig(
    level=logging.INFO if IS_MAIN else logging.WARNING,
    format=f'%(asctime)s - [Rank {RANK}] - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

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_2600.pt', 
        'data_path': '/root/dataset/r1_zero_math.jsonl',
        'save_dir': '/root/checkpoints/r1_zero_reproduction',
        'resume_from': None, 
        'model_dim': 1536,
        'n_layers': 12,
        'n_heads': 12,
        'n_kv_heads': 4,
        'group_size': 4,      
        'batch_size': 1,      
        'learning_rate': 2e-6,
        'max_steps': 190000,
        'max_gen_len': 512,
        'save_interval': 300,
        'gradient_accumulation_steps': 8, 
        'inner_batch_size': 4 
    }
    # ---------------------------------------------
    
    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"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"Configuration: {json.dumps(CONFIG, indent=2)}")

    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

    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}")
    
    logger.info("Initializing Actor Model...")
    actor = create_model().to(device)
    
    logger.info("Initializing Ref Model...")
    ref = create_model().to(device)
    ref.eval()
    ref.requires_grad_(False)
    
    trainer = GRPOZeroTrainer(
        actor_model=actor,
        ref_model=ref,
        tokenizer=tokenizer,
        learning_rate=CONFIG['learning_rate'],
        group_size=CONFIG['group_size'],
        use_amp=True,
        gradient_accumulation_steps=CONFIG['gradient_accumulation_steps'],
        inner_batch_size=CONFIG['inner_batch_size']
    )
    
    start_step = 0
    samples_seen = 0
    
    if CONFIG['resume_from']:
        resume_path = CONFIG['resume_from']
        logger.info(f"Resuming from: {resume_path}")
        checkpoint = torch.load(resume_path, map_location='cpu')
        
        actor.load_state_dict(checkpoint['model_state_dict'])
        if 'optimizer_state_dict' in checkpoint:
            try:
                trainer.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
            except Exception as e:
                logger.warning(f"Optimizer load failed (param mismatch?): {e}")
        
        ref.load_state_dict(checkpoint['model_state_dict']) 
        
        start_step = checkpoint.get('step', 0) + 1
        samples_seen = checkpoint.get('samples_seen', start_step * CONFIG['batch_size'] * WORLD_SIZE)
        
        del checkpoint
        gc.collect()
        torch.cuda.empty_cache()
    else:
        logger.info(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()}
        
        actor.load_state_dict(new_state_dict)
        ref.load_state_dict(new_state_dict)
        del checkpoint, state_dict, new_state_dict
        gc.collect()
        torch.cuda.empty_cache()
    
    if WORLD_SIZE > 1:
        actor = DDP(actor, device_ids=[LOCAL_RANK], output_device=LOCAL_RANK)
        
    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 sampler:
        epoch = samples_seen // len(dataset)
        sampler.set_epoch(epoch)
    
    data_iter = iter(dataloader)

    if samples_seen > 0:
        skip_batches = samples_seen // (CONFIG['batch_size'] * WORLD_SIZE)
        logger.info(f"Skipping {skip_batches} batches...")
        for _ in range(skip_batches):
            try:
                next(data_iter)
            except StopIteration:
                if sampler: sampler.set_epoch(sampler.epoch + 1)
                data_iter = iter(dataloader)
                next(data_iter)

    progress_bar = tqdm(range(start_step, CONFIG['max_steps']), disable=not IS_MAIN, initial=start_step, total=CONFIG['max_steps'])

    current_samples = samples_seen
    running_reward = 0.0
    running_loss = 0.0
    
    for step in progress_bar:
        try:
            try:
                batch = next(data_iter)
            except StopIteration:
                if sampler:
                    epoch = current_samples // len(dataset)
                    sampler.set_epoch(epoch)
                data_iter = iter(dataloader)
                batch = next(data_iter)
            
            current_samples += CONFIG['batch_size'] * WORLD_SIZE

            experience = trainer.generate_and_score(
                batch, 
                max_gen_len=CONFIG['max_gen_len']
            )

            step_reward = experience['avg_reward']
            if running_reward == 0: running_reward = step_reward
            else: running_reward = 0.95 * running_reward + 0.05 * step_reward

            loss = trainer.train_step(experience)
            status_dict = {"R": f"{running_reward:.3f}"}
            
            if loss is not None:
                if running_loss == 0: running_loss = loss
                else: running_loss = 0.9 * running_loss + 0.1 * loss
                status_dict["L"] = f"{running_loss:.3f}"
                
                if IS_MAIN:
                    current_lr = trainer.optimizer.param_groups[0]['lr']
                    metrics_data = {
                        "step": step,
                        "reward": float(step_reward), # 当前步的 reward
                        "loss": float(loss),
                        "lr": float(current_lr),
                        "samples_seen": current_samples,
                        "timestamp": datetime.now().isoformat()
                    }
                    with open(os.path.join(CONFIG['save_dir'], "metrics.jsonl"), "a") as f:
                        f.write(json.dumps(metrics_data) + "\n")
                        
                    if step % 10 == 0:
                        logger.info(f"Step {step} | Reward: {step_reward:.4f} | Loss: {loss:.4f} | LR: {current_lr:.2e}")
            else:
                status_dict["State"] = "Acc"
            
            progress_bar.set_description(f"{' '.join([f'{k}:{v}' for k,v in status_dict.items()])}")
            
            if step > 0 and step % CONFIG['save_interval'] == 0 and IS_MAIN:
                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': current_samples,
                    'model_state_dict': model_to_save.state_dict(),
                    'optimizer_state_dict': trainer.optimizer.state_dict(),
                }, save_path)
                logger.info(f"Checkpoint saved: {save_path}")
            
            # 显存清理
            del experience
            del batch

            
        except Exception as e:
            logger.error(f"Step {step} Error: {e}")
            import traceback
            traceback.print_exc()
            continue

    if IS_MAIN:
        final_path = f"{CONFIG['save_dir']}/final_r1_zero.pt"
        model_to_save = actor.module if hasattr(actor, 'module') else actor
        torch.save({
            'step': CONFIG['max_steps'],
            'model_state_dict': model_to_save.state_dict(),
        }, final_path)
        logger.info("Training Finished.")
        
    if WORLD_SIZE > 1:
        dist.destroy_process_group()

if __name__ == "__main__":
    main()