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