train.py

# =============================================================================
# Step 1: Import Dependencies and Dataset
# =============================================================================
import os
import math
from contextlib import nullcontext
from dataclasses import dataclass

import numpy as np
import torch
import torch.nn as nn
from torch.nn import functional as F
from torch.optim.lr_scheduler import LinearLR, SequentialLR, CosineAnnealingLR
from datasets import load_dataset
import tiktoken
from tqdm.auto import tqdm
import matplotlib.pyplot as plt

# Attempt to import Google Colab specifics, fail gracefully if not available
try:
    from google.colab import runtime
except ImportError:
    runtime = None

# =============================================================================
# Step 2: Tokenize the Dataset
# =============================================================================
# Download and load the TinyStories dataset
print("Step 1/10: Loading dataset...")
ds = load_dataset("roneneldan/TinyStories")

# Initialize the GPT-2 tokenizer
enc = tiktoken.get_encoding("gpt2")

def process(example):
    """Tokenize a single text example."""
    ids = enc.encode_ordinary(example['text']) # encode_ordinary ignores any special tokens
    out = {'ids': ids, 'len': len(ids)}
    return out

# Check if tokenized data already exists
if not os.path.exists("train.bin"):
    print("Step 2/10: Tokenizing and saving the dataset to binary files...")
    tokenized = ds.map(
        process,
        remove_columns=['text'],
        desc="Tokenizing the splits",
        num_proc=8,
    )

    # Concatenate all token IDs into large binary files for training and validation
    for split, dset in tokenized.items():
        arr_len = np.sum(dset['len'], dtype=np.uint64)
        filename = f'{split}.bin'
        # Use uint16 since vocab size (50257) < 65536
        dtype = np.uint16
        arr = np.memmap(filename, dtype=dtype, mode='w+', shape=(arr_len,))
        total_batches = 1024
        idx = 0
        for batch_idx in tqdm(range(total_batches), desc=f'Writing {filename}'):
            batch = dset.shard(num_shards=total_batches, index=batch_idx, contiguous=True).with_format('numpy')
            arr_batch = np.concatenate(batch['ids'])
            arr[idx : idx + len(arr_batch)] = arr_batch
            idx += len(arr_batch)
        arr.flush()
else:
    print("Step 2/10: Tokenized data found. Skipping tokenization.")

# =============================================================================
# Step 3: Define Training Configuration
# =============================================================================
print("Step 3/10: Defining training configuration...")
# Training Hyperparameters
learning_rate = 1e-4
max_iters = 20000
warmup_steps = 1000
min_lr = 5e-4
eval_iters = 500
batch_size = 32
block_size = 128  # Context window size
gradient_accumulation_steps = 32

# System Configuration
device = "cuda" if torch.cuda.is_available() else "cpu"
device_type = 'cuda' if 'cuda' in device else 'cpu'
dtype = 'bfloat16' if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else 'float16'
ptdtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}[dtype]
ctx = nullcontext() if device_type == 'cpu' else torch.amp.autocast(device_type=device_type, dtype=ptdtype)

torch.set_default_device(device)
torch.manual_seed(42)

# =============================================================================
# Step 4: Create Input-Output Batches
# =============================================================================
def get_batch(split):
    """
    Generate a batch of data for training or validation.
    """
    if split == 'train':
        data = np.memmap('train.bin', dtype=np.uint16, mode='r')
    else:
        data = np.memmap('validation.bin', dtype=np.uint16, mode='r')
    
    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])

    if device_type == 'cuda':
        # Pin memory for faster GPU transfer
        x, y = x.pin_memory().to(device, non_blocking=True), y.pin_memory().to(device, non_blocking=True)
    else:
        x, y = x.to(device), y.to(device)
    return x, y

# =============================================================================
# Step 5: Define the SLM Model Architecture
# =============================================================================
print("Step 5/10: Defining the GPT model architecture...")

@dataclass
class GPTConfig:
    block_size: int
    vocab_size: int
    n_layer: int
    n_head: int
    n_embd: int
    dropout: float = 0.0
    bias: bool = True

class LayerNorm(nn.Module):
    def __init__(self, ndim, bias):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(ndim))
        self.bias = nn.Parameter(torch.zeros(ndim)) if bias else None
    def forward(self, x):
        return F.layer_norm(x, self.weight.shape, self.weight, self.bias, 1e-5)

class CausalSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        assert config.n_embd % config.n_head == 0
        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
        self.attn_dropout = nn.Dropout(config.dropout)
        self.resid_dropout = nn.Dropout(config.dropout)
        self.n_head = config.n_head
        self.n_embd = config.n_embd
        self.flash = hasattr(F, 'scaled_dot_product_attention')
        if not self.flash:
            self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size))
                                        .view(1, 1, config.block_size, config.block_size))

    def forward(self, x):
        B, T, C = x.size()
        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)

        if self.flash:
            y = F.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=self.attn_dropout.p if self.training else 0.0, is_causal=True)
        else:
            att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
            att = att.masked_fill(self.bias[:, :, :T, :T] == 0, float('-inf'))
            att = F.softmax(att, dim=-1)
            att = self.attn_dropout(att)
            y = att @ v

        y = y.transpose(1, 2).contiguous().view(B, T, C)
        y = self.resid_dropout(self.c_proj(y))
        return y

class MLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias)
        self.gelu = nn.GELU()
        self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias)
        self.dropout = nn.Dropout(config.dropout)
    def forward(self, x):
        return self.dropout(self.c_proj(self.gelu(self.c_fc(x))))

class Block(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.ln1 = LayerNorm(config.n_embd, config.bias)
        self.attn = CausalSelfAttention(config)
        self.ln2 = LayerNorm(config.n_embd, config.bias)
        self.mlp = MLP(config)
    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.mlp(self.ln2(x))
        return x

class GPT(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.transformer = nn.ModuleDict(dict(
            wte=nn.Embedding(config.vocab_size, config.n_embd),
            wpe=nn.Embedding(config.block_size, config.n_embd),
            drop=nn.Dropout(config.dropout),
            h=nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
            ln_f=LayerNorm(config.n_embd, config.bias),
        ))
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
        self.transformer.wte.weight = self.lm_head.weight # weight tying

        self.apply(self._init_weights)
        for pn, p in self.named_parameters():
            if pn.endswith('c_proj.weight'):
                nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * config.n_layer))

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, idx, targets=None):
        device = idx.device
        b, t = idx.size()
        assert t <= self.config.block_size
        pos = torch.arange(0, t, dtype=torch.long, device=device)

        tok_emb = self.transformer.wte(idx)
        pos_emb = self.transformer.wpe(pos)
        x = self.transformer.drop(tok_emb + pos_emb)
        for block in self.transformer.h:
            x = block(x)
        x = self.transformer.ln_f(x)

        if targets is not None:
            logits = self.lm_head(x)
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
            return logits, loss
        else:
            logits = self.lm_head(x[:, [-1], :])
            return logits, None

    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
        """Generate tokens given a conditioning sequence."""
        for _ in range(max_new_tokens):
            idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:]
            logits, _ = self(idx_cond)
            logits = logits[:, -1, :] / temperature
            if top_k is not None:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = -float('Inf')
            probs = F.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, idx_next), dim=1)
        return idx

# Instantiate the model
config = GPTConfig(
    vocab_size=50257,
    block_size=128,
    n_layer=6,
    n_head=6,
    n_embd=384,
    dropout=0.1,
    bias=True
)
model = GPT(config)
model.to(device)

# =============================================================================
# Step 6: Define the Loss Estimation Function
# =============================================================================
@torch.no_grad()
def estimate_loss(model):
    """Computes the average loss over a fixed number of batches."""
    out = {}
    model.eval()
    for split in ['train', 'validation']:
        losses = torch.zeros(eval_iters)
        for k in range(eval_iters):
            X, Y = get_batch(split)
            with ctx:
                logits, loss = model(X, Y)
            losses[k] = loss.item()
        out[split] = losses.mean()
    model.train()
    return out

# =============================================================================
# Step 7: Define Optimizer and Scheduler
# =============================================================================
print("Step 7/10: Defining optimizer and learning rate scheduler...")
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, betas=(0.9, 0.95), weight_decay=0.1, eps=1e-9)
scheduler_warmup = LinearLR(optimizer, start_factor=0.01, total_iters=warmup_steps)
scheduler_decay = CosineAnnealingLR(optimizer, T_max=max_iters - warmup_steps, eta_min=min_lr)
scheduler = SequentialLR(optimizer, schedulers=[scheduler_warmup, scheduler_decay], milestones=[warmup_steps])
scaler = torch.cuda.amp.GradScaler(enabled=(dtype == 'float16'))

# =============================================================================
# Step 8: Pre-train the SLM
# =============================================================================
print("Step 8/10: Starting model pre-training...")
best_val_loss = float('inf')
best_model_params_path = "best_model_params.pt"
train_loss_list, validation_loss_list = [], []

# Main training loop
for epoch in tqdm(range(max_iters), desc="Training"):
    # Evaluate model and save checkpoints
    if epoch % eval_iters == 0 and epoch != 0:
        losses = estimate_loss(model)
        val_loss_cpu = losses['validation'].cpu().item() # Move to CPU for comparison
        print(f"\nEpoch {epoch}: train loss {losses['train']:.4f}, val loss {val_loss_cpu:.4f}")
        print(f"Current learning rate: {optimizer.param_groups[0]['lr']:.6f}")
        
        train_loss_list.append(losses['train'])
        validation_loss_list.append(losses['validation'])

        if val_loss_cpu < best_val_loss:
            best_val_loss = val_loss_cpu
            torch.save(model.state_dict(), best_model_params_path)
            print(f"New best model saved with validation loss: {best_val_loss:.4f}")

    # Forward and backward pass
    X, y = get_batch("train")

    with ctx:
        logits, loss = model(X, y)
        loss = loss / gradient_accumulation_steps
    
    scaler.scale(loss).backward()

    # Update weights
    if ((epoch + 1) % gradient_accumulation_steps == 0) or (epoch + 1 == max_iters):
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=0.5)
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad(set_to_none=True)
    
    scheduler.step()

print("Training finished.")

# =============================================================================
# Step 9: Plot the Loss Curves
# =============================================================================
print("Step 9/10: Plotting loss curves...")
train_loss_list_converted = [i.cpu().detach().numpy() for i in train_loss_list]
validation_loss_list_converted = [i.cpu().detach().numpy() for i in validation_loss_list]

plt.figure(figsize=(10, 6))
plt.plot(train_loss_list_converted, 'g', label='Train Loss')
plt.plot(validation_loss_list_converted, 'r', label='Validation Loss')
plt.xlabel(f"Steps (x{eval_iters})")
plt.ylabel("Loss")
plt.title("Training and Validation Loss")
plt.legend()
plt.grid(True)
plt.savefig("loss_curves.png")
plt.show()

# =============================================================================
# Step 10: Run SLM Inference on the Trained Model
# =============================================================================
print("\nStep 10/10: Running inference on the trained model...")

# Load the best model
model = GPT(config)
model.load_state_dict(torch.load(best_model_params_path, map_location=torch.device(device)))
model.to(device)
model.eval()

print("-" * 50)
# --- Inference Example 1 ---
sentence = "Once upon a time there was a pumpkin."
print(f"Prompt: {sentence}")
context = (torch.tensor(enc.encode_ordinary(sentence), device=device).unsqueeze(dim=0))
with torch.no_grad():
    with ctx:
        generated_tokens = model.generate(context, max_new_tokens=200, temperature=0.8, top_k=20)
print("\nGenerated Text:")
print(enc.decode(generated_tokens.squeeze().tolist()))
print("-" * 50)

# --- Inference Example 2 ---
sentence = "A little girl went to the woods"
print(f"Prompt: {sentence}")
context = (torch.tensor(enc.encode_ordinary(sentence), device=device).unsqueeze(dim=0))
with torch.no_grad():
    with ctx:
        generated_tokens = model.generate(context, max_new_tokens=200, temperature=0.8, top_k=20)
print("\nGenerated Text:")
print(enc.decode(generated_tokens.squeeze().tolist()))
print("-" * 50)

# Unassign Colab runtime if applicable
if runtime:
    print("Unassigning Colab runtime...")
    runtime.unassign()