Update train_baseline.py

train_baseline.py

@@ -6,18 +6,19 @@ Ensures proper bpc calculation and comparable evaluation with DTAT.
 import os
 import time
 import math
+import wandb
 import numpy as np
+from tqdm import tqdm
+
 import torch
 import torch.nn.functional as F
-from torch.nn.parallel import DistributedDataParallel as DDP
-from torch.distributed import init_process_group, destroy_process_group
-from contextlib import nullcontext
-import wandb
-from tqdm import tqdm
+from torch.nn import CrossEntropyLoss
 
 from model_baseline import BaselineTransformer
 from config.baseline_config import get_config
 
+# -----------------------------------------------------------------------------
+# I/O
 def get_batch(data, block_size, batch_size, device):
     """Generate a small batch of data of inputs x and targets y."""
     ix = torch.randint(len(data) - block_size, (batch_size,))
@@ -29,42 +30,61 @@ def get_batch(data, block_size, batch_size, device):
 def estimate_loss(model, data, config):
     """Estimate loss on data split, ensuring proper bpc calculation."""
     model.eval()
-    losses = torch.zeros(config.eval_iters)
-    for k in range(config.eval_iters):
-        X, Y = get_batch(data, config.block_size, config.batch_size, config.device)
-        with torch.no_grad():
-            logits, loss = model(X, Y)
-            losses[k] = loss.item()  # Loss is already in BPC
-    out = losses.mean()
+    total_loss = 0.0
+    total_steps = config.eval_iters
+    
+    with torch.no_grad():
+        for _ in range(total_steps):
+            X, Y = get_batch(data, config.block_size, config.batch_size, config.device)
+            with torch.amp.autocast('cuda', enabled=config.mixed_precision):
+                logits, loss = model(X, Y)
+            total_loss += loss.item()
+    
     model.train()
-    return out
+    return total_loss / total_steps
 
 def get_lr(it, config):
-    """Get learning rate based on iteration."""
+    """
+    Learning rate scheduler with linear warmup and cosine decay.
+    Matches DTAT's scheduler exactly.
+    """
+    # Linear warmup
     if it < config.warmup_iters:
         return config.learning_rate * it / config.warmup_iters
-    if it > config.lr_decay_iters:
-        return config.min_lr
-    decay_ratio = (it - config.warmup_iters) / (config.lr_decay_iters - config.warmup_iters)
-    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
-    return config.min_lr + coeff * (config.learning_rate - config.min_lr)
+    
+    # Cosine decay
+    if config.decay_lr:
+        decay_ratio = (it - config.warmup_iters) / (config.lr_decay_iters - config.warmup_iters)
+        decay_ratio = min(decay_ratio, 1.0)  # Cap at 1.0
+        coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
+        return config.min_lr + coeff * (config.learning_rate - config.min_lr)
+    
+    return config.learning_rate
 
 def main():
-    # Initialize config
-    config = get_config()
-    
     # Initialize wandb
-    wandb.init(project='enwik8-baseline', config=vars(config))
+    wandb.init(project="enwik8-baseline", name="baseline-run")
+    wandb.config.update(get_config().__dict__)
+    
+    # Get config and setup
+    config = get_config()
+    device = config.device
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+    torch.backends.cudnn.benchmark = config.cudnn_benchmark
     
-    # Load dataset
-    data = np.memmap('data/train.bin', dtype=np.uint8, mode='r')
-    val_data = np.memmap('data/val.bin', dtype=np.uint8, mode='r')
+    # Data loading
+    print("Loading data...")
+    data_dir = os.path.join('data')
+    train_data = np.memmap(os.path.join(data_dir, 'train.bin'), dtype=np.uint8, mode='r')
+    val_data = np.memmap(os.path.join(data_dir, 'val.bin'), dtype=np.uint8, mode='r')
     
-    # Initialize model
-    model = BaselineTransformer(config)
-    model.to(config.device)
+    # Model init
+    print("Initializing model...")
+    model = BaselineTransformer(config).to(device)
+    print(f"number of parameters: {model.get_num_params()/1e6:.2f}M")
     
-    # Initialize optimizer
+    # Optimizer
     optimizer = torch.optim.AdamW(
         model.parameters(),
         lr=config.learning_rate,
@@ -72,45 +92,47 @@ def main():
         weight_decay=config.weight_decay
     )
     
-    if config.compile:
-        print("Compiling model...")
-        model = torch.compile(model)
+    # Mixed precision setup
+    scaler = torch.amp.GradScaler('cuda', enabled=config.mixed_precision)
     
-    # Enable mixed precision training
-    scaler = torch.cuda.amp.GradScaler(enabled=config.mixed_precision)
-    
-    # Enable cuDNN benchmarking
-    torch.backends.cudnn.benchmark = True
+    # Memory optimizations
+    if config.gradient_checkpointing:
+        model.gradient_checkpointing_enable()
     
     # Calculate total steps and epochs
     total_steps = config.max_iters
     batch_size = config.batch_size
     block_size = config.block_size
-    total_epochs = (total_steps * batch_size * block_size) // len(data)
-    
-    print(f"Training baseline model for {total_epochs} epochs ({total_steps} iterations)")
+    total_epochs = (total_steps * batch_size * block_size) // len(train_data)
     
     # Create progress bar
     pbar = tqdm(range(config.max_iters), desc=f"Training (0/{total_epochs} epochs)")
     
     best_val_loss = float('inf')
+    no_improvement = 0
     t0 = time.time()
     
     for iter_num in pbar:
+        # Early stopping check
+        if no_improvement >= config.patience:
+            print(f"\nEarly stopping triggered after {iter_num} iterations")
+            print(f"Best validation loss: {best_val_loss:.4f}")
+            break
+        
         # Update learning rate
         lr = get_lr(iter_num, config)
         for param_group in optimizer.param_groups:
             param_group['lr'] = lr
         
-        # Sample batch
-        X, Y = get_batch(data, config.block_size, config.batch_size, config.device)
+        # Sample a batch of data
+        X, Y = get_batch(train_data, config.block_size, config.batch_size, device)
         
-        # Forward pass with mixed precision
-        with torch.cuda.amp.autocast(enabled=config.mixed_precision):
+        # Mixed precision training
+        with torch.amp.autocast('cuda', enabled=config.mixed_precision):
             logits, loss = model(X, Y)
         
         # Backward pass with gradient scaling
-        optimizer.zero_grad(set_to_none=True)
+        optimizer.zero_grad(set_to_none=True)  # Slightly faster than zero_grad()
         scaler.scale(loss).backward()
         scaler.unscale_(optimizer)
         torch.nn.utils.clip_grad_norm_(model.parameters(), config.grad_clip)
@@ -121,12 +143,17 @@ def main():
         if iter_num % config.log_interval == 0:
             # Calculate current epoch
             current_tokens = (iter_num + 1) * batch_size * block_size
-            current_epoch = current_tokens / len(data)
+            current_epoch = current_tokens / len(train_data)
             
+
+            val_loss = estimate_loss(model, val_data, config)
+
+
             # Update progress bar
             pbar.set_description(
                 f"Training ({current_epoch:.1f}/{total_epochs} epochs) | "
-                f"loss: {loss.item():.4f} | "  # Already in BPC
+                f"loss: {loss.item():.4f} | "  # This is now directly in BPC
+                f"bpc: {loss.item():.2f} | "   # Same as loss since it's already BPC
                 f"lr: {lr:.1e} | "
                 f"tokens/sec: {(batch_size * block_size) / (time.time() - t0):.1f}"
             )
@@ -134,37 +161,32 @@ def main():
             # Log to wandb
             wandb.log({
                 "iter": iter_num,
-                "loss": loss.item(),
-                "bpc": loss.item(),  # Already in BPC
-                "lr": lr,
                 "epoch": current_epoch,
+                "train/loss": loss.item(),
+                "train/bpc": loss.item(),  # Same as loss since it's already BPC
+                "lr": lr,
                 "tokens_per_sec": (batch_size * block_size) / (time.time() - t0),
             })
-            
-            t0 = time.time()
         
-        # Evaluation
-        if iter_num > 0 and iter_num % config.eval_interval == 0:
+        # Check validation and save every 100 iterations
+        if iter_num > 0 and iter_num % 100 == 0:
             val_loss = estimate_loss(model, val_data, config)
-            wandb.log({
-                "val_loss": val_loss,
-                "val_bpc": val_loss,  # Already in BPC
-                "epoch": current_epoch,
-            })
-            
-            # Save best model
             if val_loss < best_val_loss:
                 best_val_loss = val_loss
-                print(f"Saving best model with val_bpc: {val_loss:.4f}")
-                torch.save(model.state_dict(), 'models/baseline_best.pt')
-    
-    # Final evaluation
-    model.eval()
-    final_val_loss = estimate_loss(model, val_data, config)
-    print(f"Final validation BPC: {final_val_loss:.4f}")
+                no_improvement = 0
+                print(f"\nSaving best model with val_loss: {best_val_loss:.4f}")
+                torch.save(model.state_dict(), os.path.join(os.path.dirname(__file__), 'best_baseline.pt'))
+            else:
+                no_improvement += 1
+            
+            # Log validation loss to wandb
+            wandb.log({
+                "val/loss": val_loss,
+                "val/bpc": val_loss,
+                "lr": lr,
+
+            })
     
-    # Save final model
-    torch.save(model.state_dict(), 'models/baseline_final.pt')
     wandb.finish()
 
 if __name__ == '__main__':