gpt_training.py

import os, time
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.amp import GradScaler, autocast  # new API

# --------- Setup ----------
device = "cuda" if torch.cuda.is_available() else "cpu"
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True

# Use FP16 autocast + grad scaling on CUDA
scaler = GradScaler(enabled=(device == "cuda"))

print("==== Runtime Info ====")
print("PyTorch:", torch.__version__)
print("Device:", device, (torch.cuda.get_device_name(0) if device=="cuda" else ""))
print("TF32 matmul:", torch.backends.cuda.matmul.allow_tf32)
print("TF32 cudnn :", torch.backends.cudnn.allow_tf32)
print("======================")

# --------- Load dataset ----------
with open("input.txt", "r", encoding="utf-8") as file:
    text = file.read()
print(f"Loaded dataset: {len(text):,} chars")

# --------- Tokenizer ----------
vocab = sorted(list(set(text)))
vocab_size = len(vocab)
stoi = {c: i for i, c in enumerate(vocab)}
itos = {i: c for i, c in enumerate(vocab)}
encode = lambda s: [stoi[c] for c in s]
decode = lambda l: ''.join([itos[t] for t in l])
print(f"Vocab size: {vocab_size}")

# --------- Dataset & batching ----------
dataset = torch.tensor(encode(text), dtype=torch.long)
split_point = int(0.9 * len(dataset))
train_data = dataset[:split_point]
val_data = dataset[split_point:]
print(f"Train tokens: {len(train_data):,} | Val tokens: {len(val_data):,}")

def get_batch(split, batch_size, block_size):
    data = train_data if split == "train" else val_data
    ix = torch.randint(0, len(data) - block_size - 1, (batch_size,))
    x = torch.stack([data[i:i+block_size] for i in ix])
    y = torch.stack([data[i+1:i+1+block_size] for i in ix])
    return x.to(device), y.to(device)

# --------- Hyperparameters ----------
torch.manual_seed(1337)
batch_size    = 64
block_size    = 256
max_iters     = 5000
eval_interval = 50
learning_rate = 3e-4
eval_iters    = 200
n_embd        = 384
n_layer       = 6
n_heads       = 6
dropout       = 0.2

# --------- Model ----------
class Head(nn.Module):
    def __init__(self, head_size):
        super().__init__()
        self.query = nn.Linear(n_embd, head_size, bias=False)
        self.key   = nn.Linear(n_embd, head_size, bias=False)
        self.value = nn.Linear(n_embd, head_size, bias=False)
        self.register_buffer("tril", torch.tril(torch.ones(block_size, block_size)))
        self.dropout = nn.Dropout(dropout)
        self.scale = head_size ** -0.5
    def forward(self, x):
        B, T, C = x.shape
        k = self.key(x); q = self.query(x); v = self.value(x)
        wei = (q @ k.transpose(-2, -1)) * self.scale
        wei = wei.masked_fill(self.tril[:T, :T] == 0, float("-inf"))
        wei = F.softmax(wei, dim=-1)
        wei = self.dropout(wei)
        return wei @ v

class MultiHeadAttention(nn.Module):
    def __init__(self, num_heads, head_size):
        super().__init__()
        self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])
        self.proj = nn.Linear(head_size * num_heads, n_embd)
        self.dropout = nn.Dropout(dropout)
    def forward(self, x):
        out = torch.cat([h(x) for h in self.heads], dim=-1)
        return self.dropout(self.proj(out))

class FeedForward(nn.Module):
    def __init__(self, n_embd):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(n_embd, 4 * n_embd),
            nn.GELU(),
            nn.Linear(4 * n_embd, n_embd),
            nn.Dropout(dropout),
        )
    def forward(self, x): return self.net(x)

class Block(nn.Module):
    def __init__(self, n_embd, num_heads):
        super().__init__()
        head_size = n_embd // num_heads
        self.sa = MultiHeadAttention(num_heads, head_size)
        self.ffwd = FeedForward(n_embd)
        self.ln1 = nn.LayerNorm(n_embd)
        self.ln2 = nn.LayerNorm(n_embd)
    def forward(self, x):
        x = x + self.sa(self.ln1(x))
        x = x + self.ffwd(self.ln2(x))
        return x

class GPT(nn.Module):
    def __init__(self):
        super().__init__()
        self.token_embedding_table = nn.Embedding(vocab_size, n_embd)
        self.position_embedding_table = nn.Embedding(block_size, n_embd)
        self.blocks = nn.Sequential(*[Block(n_embd, n_heads) for _ in range(n_layer)])
        self.ln_f = nn.LayerNorm(n_embd)
        self.lm_head = nn.Linear(n_embd, vocab_size)
    def forward(self, idx, targets=None):
        B, T = idx.shape
        tok_emb = self.token_embedding_table(idx)
        pos_emb = self.position_embedding_table(torch.arange(0, T, device=idx.device))
        x = tok_emb + pos_emb
        x = self.blocks(x)
        x = self.ln_f(x)
        logits = self.lm_head(x)
        loss = None
        if targets is not None:
            B, T, C = logits.shape
            loss = F.cross_entropy(logits.view(B*T, C), targets.view(B*T))
        return logits, loss
    @torch.no_grad()
    def generate(self, idx, max_new_tokens):
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -block_size:]
            logits, _ = self(idx_cond)
            logits = logits[:, -1, :]
            probs = F.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, idx_next), dim=1)
        return idx

# --------- Eval helper ----------
@torch.no_grad()
def estimate_loss(model: GPT):
    out = {}
    model.eval()
    for split in ["train", "val"]:
        losses = torch.zeros(eval_iters, device=device)
        for i in range(eval_iters):                         
            X, Y = get_batch(split, batch_size, block_size)
            with autocast(device_type="cuda", dtype=torch.float16, enabled=(device=="cuda")):
                logits, loss = model(X, Y)
            losses[i] = loss.item()
        out[split] = losses.mean().item()
    model.train()
    return out

# --------- Train ----------
model = GPT().to(device)
n_params = sum(p.numel() for p in model.parameters())
print(f"Model params: {n_params/1e6:.2f}M")

optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)

start = time.time()
for it in range(max_iters):
    if it % eval_interval == 0:
        losses = estimate_loss(model)
        elapsed = time.time() - start
        print(f"[{it:5d}/{max_iters}] train {losses['train']:.4f} | val {losses['val']:.4f} | elapsed {elapsed:.1f}s")

    xb, yb = get_batch("train", batch_size, block_size)
    optimizer.zero_grad(set_to_none=True)
    with autocast(device_type="cuda", dtype=torch.float16, enabled=(device=="cuda")):
        _, loss = model(xb, yb)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

print("\nTraining complete.")
os.makedirs("saved_models", exist_ok=True)
save_path = os.path.join("saved_models", "model_params.pth")
torch.save(model.state_dict(), save_path)
print("Saved model to:", save_path)