Update pretrain.py

pretrain.py

@@ -55,26 +55,21 @@ class PreTrainer:
         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
@@ -89,17 +84,13 @@ class PreTrainer:
         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
 
@@ -115,9 +106,6 @@ class PreTrainer:
                 positions = torch.cumsum(non_pad_mask.long(), dim=0) -1
                 position_ids[i]=positions * non_pad_mask.long()
 
-
-        
-        # 准备输入
         input_data = {
             'segments': [{
                 'type': 'text',
@@ -126,7 +114,6 @@ class PreTrainer:
             }]
         }
         
-        # 前向传播
         with torch.amp.autocast('cuda', enabled=self.use_amp):
             outputs = self.model(
                 input_data,
@@ -134,7 +121,6 @@ class PreTrainer:
                 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()
@@ -145,7 +131,6 @@ class PreTrainer:
                 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
         
@@ -153,27 +138,22 @@ class PreTrainer:
         self.tokens_seen += attention_mask.sum().item()
         
         return {
-            'loss': loss.item(),  # 返回真实的、未缩放的loss
+            'loss': loss.item(), 
             '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()
@@ -182,7 +162,6 @@ class PreTrainer:
         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}, "
@@ -194,19 +173,11 @@ class PreTrainer:
             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("Training Log (Real Loss Values)\n")
                 f.write("="*80 + "\n")
         
         self.model.train()
@@ -219,26 +190,21 @@ class PreTrainer:
         
         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...")
+            logger.info(f"Resuming: Need to skip {batches_to_skip} batches to restore data state...")
         
         # 创建迭代器
         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.")
+                        logger.error("Dataset exhausted during skipping!")
                         return
-        
-        logger.info(" Data fast-forward complete. Resuming training...")
 
         try:
             while True:
@@ -251,17 +217,20 @@ class PreTrainer:
                     continue
                 stats = self.train_step(batch)
                 step_in_accumulation += 1
-                accumulated_loss += stats['loss'] 
+                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 
+                    stats['loss'] = avg_step_loss  
+                    
                     self.running_loss += avg_step_loss
                     
                     step_in_accumulation = 0
-                    accumulated_loss = 0.0 
+                    accumulated_loss = 0.0  
+                    
                     progress_bar.update(1)
                     progress_bar.set_postfix({
                         'loss': f"{stats['loss']:.4f}",
@@ -270,7 +239,6 @@ class PreTrainer:
                         'grad': f"{grad_norm:.2f}"
                     })
                     
-                    # 日志记录
                     if self.global_step % self.log_interval == 0:
                         avg_loss = self.running_loss / self.log_interval
                         
@@ -293,37 +261,26 @@ class PreTrainer:
                             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:
+            logger.info("\n  Training interrupted by user")
             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(),
@@ -339,7 +296,6 @@ class PreTrainer:
         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'])
@@ -352,49 +308,38 @@ class PreTrainer:
         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,  
+        'max_seq_len': 1024,  
         'dropout': 0.1,
         'use_moe': False,
+        
+        
         'batch_size': 4,           
         'gradient_accumulation_steps': 8,  
-        'learning_rate': 3e-4,    
+        'learning_rate': 1e-4,     
         'weight_decay': 0.1,
         'warmup_steps': 500,       
-        'max_steps': 10000,
+        'max_steps': 100000,
         'max_grad_norm': 1.0,
         
-        # 数据配置
-        'data_mix': 'text_only',
-        'max_length': 512,        
+        'data_mix': 'skypile_training',
+        'max_length': 1024,        
         'num_workers': 2,          
         
-        # 日志和保存
         'log_interval': 10,
-        'save_interval': 500,      
+        'save_interval': 5000,      
         'checkpoint_dir': 'checkpoints/pretrain_fixed',
-        'loss_log_file': 'checkpoints/pretrain_fixed/train_loss.log'
+        'loss_log_file': 'checkpoints/pretrain_fixed/train_loss_skypile_training.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,
@@ -406,9 +351,7 @@ def main():
         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'],
@@ -424,9 +367,7 @@ def main():
         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,
@@ -434,8 +375,6 @@ def main():
         num_workers=config['num_workers'],
         max_length=config['max_length']
     )
-    
-    # 创建训练器
     trainer = PreTrainer(
         model=model,
         tokenizer=tokenizer,
@@ -450,9 +389,9 @@ def main():
         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, resume_from="/root/checkpoints/pretrain_fixed/step_35000.pt")
     #trainer.train(dataloader)