import os import torch import torch.nn.functional as F from transformers import AutoTokenizer from pathlib import Path import logging from tqdm import tqdm import json from datetime import datetime from model import MultiModalDenseTransformer from data_loader import create_pretrain_dataloader logging.basicConfig( level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' ) logger = logging.getLogger(__name__) os.environ["HF_ENDPOINT"] = "https://hf-mirror.com" class PreTrainer: def __init__( self, model: MultiModalDenseTransformer, tokenizer, learning_rate: float = 3e-4, weight_decay: float = 0.1, warmup_steps: int = 1000, max_steps: int = 100000, gradient_accumulation_steps: int = 16, max_grad_norm: float = 1.0, log_interval: int = 10, save_interval: int = 1000, checkpoint_dir: str = "checkpoints/pretrain", loss_log_file: str = "checkpoints/pretrain/train_loss.log" ): self.model = model 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(), lr=learning_rate, weight_decay=weight_decay, betas=(0.9, 0.95), eps=1e-8 ) from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts self.warmup_steps = warmup_steps self.max_lr = learning_rate self.min_lr = learning_rate * 0.1 self.current_step = 0 # 混合精度 self.use_amp = torch.cuda.is_available() self.scaler = torch.amp.GradScaler('cuda', enabled=self.use_amp) # 训练参数 self.gradient_accumulation_steps = gradient_accumulation_steps self.max_grad_norm = max_grad_norm self.max_steps = max_steps self.log_interval = log_interval self.save_interval = save_interval # Checkpoint管理 self.checkpoint_dir = Path(checkpoint_dir) self.checkpoint_dir.mkdir(parents=True, exist_ok=True) # 损失日志 self.loss_log_file = Path(loss_log_file) self.loss_log_file.parent.mkdir(parents=True, exist_ok=True) # 训练状态 self.global_step = 0 self.tokens_seen = 0 self.running_loss = 0.0 self.best_loss = float('inf') logger.info(f"PreTrainer initialized:") logger.info(f" Device: {self.device}") logger.info(f" Learning Rate: {learning_rate}") logger.info(f" Max Steps: {max_steps}") logger.info(f" Gradient Accumulation: {gradient_accumulation_steps}") logger.info(f" Effective Batch Size: {gradient_accumulation_steps}") logger.info(f" Mixed Precision: {self.use_amp}") def _get_lr(self) -> float: """手动计算学习率（Warmup + Cosine）""" if self.current_step < self.warmup_steps: # Linear warmup return self.max_lr * (self.current_step / self.warmup_steps) else: # Cosine decay progress = (self.current_step - self.warmup_steps) / (self.max_steps - self.warmup_steps) return self.min_lr + (self.max_lr - self.min_lr) * 0.5 * (1 + torch.cos(torch.tensor(progress * 3.14159))) def _set_lr(self, lr: float): """设置学习率""" for param_group in self.optimizer.param_groups: param_group['lr'] = lr def train_step(self, batch: dict) -> dict: input_ids = batch['input_ids'].to(self.device) attention_mask = batch['attention_mask'].to(self.device) batch_size, seq_len = 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() # 准备输入 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 = input_ids[:, 1:].contiguous() shift_attention_mask = attention_mask[:, 1:].contiguous() loss = F.cross_entropy( shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1), reduction='none' ) # 应用mask loss = (loss * shift_attention_mask.view(-1)).sum() / (shift_attention_mask.sum() + 1e-8) loss_for_backward = loss / self.gradient_accumulation_steps self.scaler.scale(loss_for_backward).backward() self.tokens_seen += attention_mask.sum().item() return { 'loss': loss.item(), # 返回真实的、未缩放的loss 'lr': self.optimizer.param_groups[0]['lr'] } def optimizer_step(self): """优化器步骤""" # Unscale梯度 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.current_step += 1 self.global_step += 1 lr = self._get_lr() self._set_lr(lr) return grad_norm.item() def _write_loss_to_txt(self, step, avg_loss, lr, tokens_seen): """写入损失日志""" log_content = ( f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] " f"Step: {step}/{self.max_steps}, " f"Average Loss: {avg_loss:.4f}, " f"Learning Rate: {lr:.2e}, " f"Tokens Seen: {tokens_seen/1e9:.2f}B\n" ) with open(self.loss_log_file, 'a', encoding='utf-8') as f: f.write(log_content) def train(self, dataloader, resume_from=None): """训练循环""" logger.info("\n" + "="*80) logger.info("Starting Pre-Training (Fixed Version)") logger.info("="*80 + "\n") # 恢复训练 if resume_from: self.load_checkpoint(resume_from) # 初始化日志 if not self.loss_log_file.exists(): with open(self.loss_log_file, 'w', encoding='utf-8') as f: f.write(" Fixed Training Log (Real Loss Values)\n") f.write("="*80 + "\n") self.model.train() progress_bar = tqdm(total=self.max_steps, initial=self.global_step) step_in_accumulation = 0 accumulated_loss = 0.0 batches_to_skip = self.global_step * self.gradient_accumulation_steps logger.info(f"Current Global Step: {self.global_step}") if batches_to_skip > 0: logger.info(f" Resuming: Need to skip {batches_to_skip} batches to restore data state...") logger.info("This might take a while depending on network/disk speed...") # 创建迭代器 data_iterator = iter(dataloader) skipped = 0 if batches_to_skip > 0: with tqdm(total=batches_to_skip, desc="Skipping trained batches", unit="batch") as skip_pbar: while skipped < batches_to_skip: try: # 只取数据，不进模型，不计算梯度 _ = next(data_iterator) skipped += 1 skip_pbar.update(1) except StopIteration: logger.error("Dataset exhausted during skipping! Check your dataset size or max_steps.") return logger.info(" Data fast-forward complete. Resuming training...") try: while True: try: batch = next(data_iterator) except StopIteration: break if batch is None or batch['input_ids'].size(0) == 0: continue stats = self.train_step(batch) step_in_accumulation += 1 accumulated_loss += stats['loss'] if step_in_accumulation >= self.gradient_accumulation_steps: avg_step_loss = accumulated_loss / self.gradient_accumulation_steps grad_norm = self.optimizer_step() stats['grad_norm'] = grad_norm stats['loss'] = avg_step_loss self.running_loss += avg_step_loss step_in_accumulation = 0 accumulated_loss = 0.0 progress_bar.update(1) progress_bar.set_postfix({ 'loss': f"{stats['loss']:.4f}", 'lr': f"{stats['lr']:.2e}", 'tokens': f"{self.tokens_seen/1e9:.2f}B", 'grad': f"{grad_norm:.2f}" }) # 日志记录 if self.global_step % self.log_interval == 0: avg_loss = self.running_loss / self.log_interval logger.info( f"Step {self.global_step}/{self.max_steps} | " f"Loss: {avg_loss:.4f} | " f"LR: {stats['lr']:.2e} | " f"GradNorm: {grad_norm:.2f} | " f"Tokens: {self.tokens_seen/1e9:.2f}B" ) if avg_loss < self.best_loss: self.best_loss = avg_loss logger.info(f" New best loss: {self.best_loss:.4f}") self._write_loss_to_txt( step=self.global_step, avg_loss=avg_loss, lr=stats['lr'], tokens_seen=self.tokens_seen ) self.running_loss = 0.0 # 保存checkpoint if self.global_step % self.save_interval == 0: self.save_checkpoint( self.checkpoint_dir / f"step_{self.global_step}.pt" ) # 完成训练 if self.global_step >= self.max_steps: break except KeyboardInterrupt: self.save_checkpoint( self.checkpoint_dir / f"interrupted_step_{self.global_step}.pt" ) finally: progress_bar.close() logger.info("\n" + "="*80) logger.info("Pre-Training Complete!") logger.info(f" Total Steps: {self.global_step}") logger.info(f" Total Tokens: {self.tokens_seen/1e9:.2f}B") logger.info(f" Best Loss: {self.best_loss:.4f}") logger.info("="*80 + "\n") # 保存最终模型 self.save_checkpoint(self.checkpoint_dir / "final_model.pt") def save_checkpoint(self, path: Path): """保存checkpoint""" checkpoint = { 'model_state_dict': self.model.state_dict(), 'optimizer_state_dict': self.optimizer.state_dict(), 'scaler_state_dict': self.scaler.state_dict() if self.use_amp else None, 'global_step': self.global_step, 'current_step': self.current_step, 'tokens_seen': self.tokens_seen, 'best_loss': self.best_loss, 'timestamp': datetime.now().isoformat() } torch.save(checkpoint, path) logger.info(f" Checkpoint saved to {path}") def load_checkpoint(self, path: str): """加载checkpoint""" checkpoint = torch.load(path, map_location=self.device, weights_only=True) self.model.load_state_dict(checkpoint['model_state_dict']) self.optimizer.load_state_dict(checkpoint['optimizer_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.current_step = checkpoint.get('current_step', self.global_step) self.tokens_seen = checkpoint['tokens_seen'] self.best_loss = checkpoint.get('best_loss', float('inf')) logger.info(f" Checkpoint loaded from {path}") logger.info(f" Resuming from step {self.global_step}") logger.info(f" Tokens seen: {self.tokens_seen/1e9:.2f}B") def main(): config = { # 模型配置 'model_dim': 1536, 'vocab_size': 151665, 'n_layers': 12, 'n_heads': 12, 'n_kv_heads': 4, 'max_seq_len': 512, 'dropout': 0.1, 'use_moe': False, 'batch_size': 4, 'gradient_accumulation_steps': 8, 'learning_rate': 3e-4, 'weight_decay': 0.1, 'warmup_steps': 500, 'max_steps': 10000, 'max_grad_norm': 1.0, # 数据配置 'data_mix': 'text_only', 'max_length': 512, 'num_workers': 2, # 日志和保存 'log_interval': 10, 'save_interval': 500, 'checkpoint_dir': 'checkpoints/pretrain_fixed', 'loss_log_file': 'checkpoints/pretrain_fixed/train_loss.log' } logger.info("="*80) logger.info(json.dumps(config, indent=2)) logger.info("="*80 + "\n") # 初始化tokenizer logger.info("Initializing 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) logger.info(f"Vocab size: {config['vocab_size']}\n") # 初始化模型 logger.info("Initializing 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=True, rope_scaling_type="yarn", use_multimodal_fusion=False, use_contrastive=False ) # 创建数据加载器 logger.info(f"\nCreating dataloader (mix: {config['data_mix']})...") dataloader = create_pretrain_dataloader( mix_name=config['data_mix'], tokenizer=tokenizer, batch_size=config['batch_size'], num_workers=config['num_workers'], max_length=config['max_length'] ) # 创建训练器 trainer = PreTrainer( model=model, tokenizer=tokenizer, learning_rate=config['learning_rate'], weight_decay=config['weight_decay'], warmup_steps=config['warmup_steps'], max_steps=config['max_steps'], gradient_accumulation_steps=config['gradient_accumulation_steps'], max_grad_norm=config['max_grad_norm'], log_interval=config['log_interval'], save_interval=config['save_interval'], checkpoint_dir=config['checkpoint_dir'], loss_log_file=config['loss_log_file'] ) logger.info("\n Starting fresh training with fixes...\n") trainer.train(dataloader, resume_from="/root/step_6500.pt") #trainer.train(dataloader) if __name__ == "__main__": main()