Refactor Training and Validation for Correct Resumption and Evaluation

Author: Pomilon

aetheris/cli/main.py

@@ -55,10 +55,17 @@ def train_command(args):
     # --- STAGE 1: PRE-TRAINING ---
     if current_stage == "Pre-Training" or start_step == 0:
         pt_loader = create_streaming_loader("cerebras/SlimPajama-627B", "train",
-                                           tokenizer, config, args.batch_size, mode="pretrain", hf_token=args.hf_token)
+                                           tokenizer, config, args.batch_size, mode="pretrain", 
+                                           hf_token=args.hf_token, start_step=start_step)
 
+        # Validation loader (no skipping needed, always from start of val set)
+        pt_val_loader = create_streaming_loader("cerebras/SlimPajama-627B", "validation",
+                                               tokenizer, config, args.batch_size, mode="pretrain", 
+                                               hf_token=args.hf_token)
+
         start_step = trainer.train_epoch(pt_loader, total_steps=args.pretrain_steps, 
-                                       start_step=start_step, stage_name="Pre-Training")
+                                       start_step=start_step, stage_name="Pre-Training",
+                                       val_loader=pt_val_loader)
         current_stage = "SFT"
         start_step = 0
 
@@ -68,10 +75,16 @@ def train_command(args):
         param_group['lr'] = 5e-5
 
     sft_loader = create_streaming_loader("OpenAssistant/oasst1", "train",
-                                        tokenizer, config, args.batch_size, mode="sft", hf_token=args.hf_token)
+                                        tokenizer, config, args.batch_size, mode="sft", 
+                                        hf_token=args.hf_token, start_step=start_step)
+
+    sft_val_loader = create_streaming_loader("OpenAssistant/oasst1", "validation",
+                                            tokenizer, config, args.batch_size, mode="sft", 
+                                            hf_token=args.hf_token)
 
     trainer.train_epoch(sft_loader, total_steps=args.sft_steps, 
-                      start_step=start_step, stage_name="SFT")
+                      start_step=start_step, stage_name="SFT",
+                      val_loader=sft_val_loader)
 
     print("\nTraining Complete!")
 

aetheris/data.py

@@ -13,12 +13,13 @@ def get_tokenizer(model_name: str = "gpt2"):
     return tokenizer
 
 class StreamingDataset(IterableDataset):
-    def __init__(self, dataset, tokenizer, max_seq_len, mode="pretrain", buffer_size=500):
+    def __init__(self, dataset, tokenizer, max_seq_len, mode="pretrain", buffer_size=500, skip_samples=0):
         self.dataset = dataset
         self.tokenizer = tokenizer
         self.max_seq_len = max_seq_len
         self.mode = mode
         self.buffer_size = buffer_size
+        self.skip_samples = skip_samples
 
     def _prepare_sft_text(self, example):
         if 'messages' in example:
@@ -39,7 +40,10 @@ def _prepare_sft_text(self, example):
     def __iter__(self) -> Iterator[Dict[str, torch.Tensor]]:
         iterator = iter(self.dataset)
         buffer = []
-
+        
+        # Calculate roughly how many items to skip if they were yielded
+        # We process skipping in the yield loop
+        
         for example in iterator:
             text = (example.get('text', '') if self.mode == "pretrain"
                    else self._prepare_sft_text(example))
@@ -70,16 +74,32 @@ def __iter__(self) -> Iterator[Dict[str, torch.Tensor]]:
             if len(buffer) >= self.buffer_size:
                 random.shuffle(buffer)
                 for _ in range(self.buffer_size // 2):
-                    yield buffer.pop()
+                    item = buffer.pop()
+                    if self.skip_samples > 0:
+                        self.skip_samples -= 1
+                        continue
+                    yield item
 
         # Yield remaining
         random.shuffle(buffer)
         while buffer:
-            yield buffer.pop()
+            item = buffer.pop()
+            if self.skip_samples > 0:
+                self.skip_samples -= 1
+                continue
+            yield item
 
-def create_streaming_loader(dataset_name, split, tokenizer, config, batch_size, mode="pretrain", hf_token=None):
+def create_streaming_loader(dataset_name, split, tokenizer, config, batch_size, mode="pretrain", hf_token=None, start_step=0):
     raw_dataset = load_dataset(dataset_name, split=split, streaming=True,
                               trust_remote_code=True, token=hf_token)
-    stream_ds = StreamingDataset(raw_dataset, tokenizer, config.max_seq_len, mode=mode)
+    
+    # Calculate samples to skip: start_step * batch_size
+    skip_samples = start_step * batch_size
+    if skip_samples > 0:
+        print(f"  [Loader] Resuming: Fast-forwarding dataset by {skip_samples} samples...")
+        
+    stream_ds = StreamingDataset(raw_dataset, tokenizer, config.max_seq_len, mode=mode, skip_samples=skip_samples)
+    
+    # Increase num_workers for better utilization
     return DataLoader(stream_ds, batch_size=batch_size, pin_memory=True,
-                     num_workers=1, prefetch_factor=2)
+                     num_workers=4, prefetch_factor=4)

aetheris/trainer/trainer.py

@@ -15,7 +15,48 @@ def __init__(self, model, optimizer, scaler, config, device, checkpoint_dir, log
 
         self.model.to(self.device)
 
-    def train_epoch(self, train_loader, total_steps, start_step=0, stage_name="Training"):
+    def validate(self, val_loader, global_step):
+        self.model.eval()
+        total_loss = 0
+        total_items = 0
+        num_batches = 100 # Validate on 100 batches to save time
+        
+        print(f"\n[Validation] Starting validation at step {global_step}...")
+        
+        with torch.no_grad():
+             for i, batch in enumerate(val_loader):
+                if i >= num_batches:
+                    break
+                    
+                input_ids, labels = batch
+                input_ids = input_ids.to(self.device, non_blocking=True)
+                labels = labels.to(self.device, non_blocking=True)
+                
+                # Auto-cast context
+                if self.device.type == 'cuda':
+                    autocast_dtype = torch.float16
+                else:
+                    autocast_dtype = torch.bfloat16
+                    
+                use_autocast = True if self.config.torch_dtype != torch.float32 else False
+                
+                if use_autocast:
+                    with torch.amp.autocast('cuda' if self.device.type == 'cuda' else 'cpu', dtype=autocast_dtype):
+                        output = self.model(input_ids, labels)
+                else:
+                    output = self.model(input_ids, labels)
+                
+                total_loss += output["loss"].item()
+                total_items += 1
+        
+        avg_loss = total_loss / total_items if total_items > 0 else 0
+        perplexity = torch.exp(torch.tensor(avg_loss)).item()
+        
+        print(f"[Validation] Step {global_step} | Loss: {avg_loss:.4f} | PPL: {perplexity:.4f}")
+        self.model.train()
+        return avg_loss
+
+    def train_epoch(self, train_loader, total_steps, start_step=0, stage_name="Training", val_loader=None, eval_every=500):
         print(f"\n{'='*70}\nStarting {stage_name}: Target Steps={total_steps}\n{'='*70}")
         self.model.train()
         global_step = start_step
@@ -29,9 +70,7 @@ def train_epoch(self, train_loader, total_steps, start_step=0, stage_name="Train
         while global_step < total_steps:
             step_start = time.time()
 
-            # Clear cache periodically
-            if global_step % 100 == 0:
-                torch.cuda.empty_cache()
+            # Removed periodic cache clearing for performance
 
             self.optimizer.zero_grad(set_to_none=True)
 
@@ -99,5 +138,8 @@ def train_epoch(self, train_loader, total_steps, start_step=0, stage_name="Train
 
             if global_step % 500 == 0:
                 save_checkpoint(self.model, self.optimizer, self.scaler, global_step, stage_name, self.checkpoint_dir)
+                
+            if val_loader and global_step % eval_every == 0 and global_step > start_step:
+                self.validate(val_loader, global_step)
 
         return global_step