train_baseline.py

"""
Training script for baseline NanoGPT model on enwik8 dataset.
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 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,))
    x = torch.stack([torch.from_numpy((data[i:i+block_size]).astype(np.int64)) for i in ix])
    y = torch.stack([torch.from_numpy((data[i+1:i+1+block_size]).astype(np.int64)) for i in ix])
    x, y = x.to(device), y.to(device)
    return x, y

def estimate_loss(model, data, config):
    """Estimate loss on data split, ensuring proper bpc calculation."""
    model.eval()
    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 total_loss / total_steps

def get_lr(it, config):
    """
    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
    
    # 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 wandb
    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
    
    # 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')
    
    # Model init
    print("Initializing model...")
    model = BaselineTransformer(config).to(device)
    print(f"number of parameters: {model.get_num_params()/1e6:.2f}M")
    
    # Optimizer
    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=config.learning_rate,
        betas=(config.beta1, config.beta2),
        weight_decay=config.weight_decay
    )
    
    # Mixed precision setup
    scaler = torch.amp.GradScaler('cuda', enabled=config.mixed_precision)
    
    # 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(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 a batch of data
        X, Y = get_batch(train_data, config.block_size, config.batch_size, device)
        
        # 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)  # Slightly faster than zero_grad()
        scaler.scale(loss).backward()
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), config.grad_clip)
        scaler.step(optimizer)
        scaler.update()
        
        # Logging
        if iter_num % config.log_interval == 0:
            # Calculate current epoch
            current_tokens = (iter_num + 1) * batch_size * block_size
            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} | "  # 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}"
            )
            
            # Log to wandb
            wandb.log({
                "iter": iter_num,
                "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),
            })
        
        # Check validation and save every 100 iterations
        if iter_num > 0 and iter_num % 100 == 0:
            val_loss = estimate_loss(model, val_data, config)
            if val_loss < best_val_loss:
                best_val_loss = val_loss
                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,

            })
    
    wandb.finish()

if __name__ == '__main__':
    main()