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