train.py

#!/usr/bin/env python3
"""

Train a small GPT (GPT-2-small scale, ~124M params by default) on the

wiki_dump.txt produced by wiki_dump_from_dataset.py, on a single RTX 4090.



What it does:

1. Tokenizes wiki_dump.txt with GPT-2 BPE (tiktoken), stripping the

   -----START PAGE----- / -----END PAGE----- markers first.

   Caches the tokenized result to .bin files next to the input so re-runs

   skip tokenization.

2. Trains a nanoGPT-style decoder-only transformer with bf16 autocast,

   flash attention (via F.scaled_dot_product_attention), and torch.compile.

3. Stops on a TIME BUDGET (default 25 min) rather than a fixed iteration

   count, since throughput varies by host — this keeps you inside a

   30-minute rental window regardless of exact hardware.

4. Shows a tqdm progress bar with live loss + ETA, and prints a sample

   generation at the end so you can sanity-check the result.



Usage:

    pip install torch tiktoken tqdm numpy

    python train_wiki_gpt.py --data wiki_dump.txt --max-minutes 25



Defaults target GPT-2-small (12 layer / 12 head / 768 dim, ~124M params)

sized to comfortably fit a 24GB 4090 with room to spare. Pass --n-layer /

--n-head / --n-embd to shrink further if you want faster iterations.

"""

import argparse
import math
import os
import re
import sys
import time

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from tqdm import tqdm

START_MARK = "-----START PAGE-----"
END_MARK = "-----END PAGE-----"


# --------------------------------------------------------------------------
# Data prep: strip markers, tokenize with GPT-2 BPE, cache to .bin
# --------------------------------------------------------------------------

def prepare_data(txt_path: str, val_fraction: float = 0.01):
    import tiktoken

    base = os.path.splitext(txt_path)[0]
    train_bin = base + ".train.bin"
    val_bin = base + ".val.bin"

    if os.path.exists(train_bin) and os.path.exists(val_bin):
        print(f"Found cached tokens: {train_bin}, {val_bin}")
        return train_bin, val_bin

    print("Tokenizing (first run only, cached after this)...")
    with open(txt_path, "r", encoding="utf-8", errors="ignore") as f:
        raw = f.read()

    # Strip the page markers and titles-as-separators; keep article text.
    raw = raw.replace(START_MARK, "").replace(END_MARK, "")

    enc = tiktoken.get_encoding("gpt2")
    ids = enc.encode_ordinary(raw)
    ids = np.array(ids, dtype=np.uint16)  # gpt2 vocab (50257) fits in uint16

    n_val = int(len(ids) * val_fraction)
    train_ids = ids[:-n_val] if n_val > 0 else ids
    val_ids = ids[-n_val:] if n_val > 0 else ids[-1000:]

    train_ids.tofile(train_bin)
    val_ids.tofile(val_bin)
    print(f"Tokenized: {len(ids):,} tokens total "
          f"({len(train_ids):,} train / {len(val_ids):,} val)")
    return train_bin, val_bin


def get_batch(bin_path: str, block_size: int, batch_size: int, device: str):
    data = np.memmap(bin_path, dtype=np.uint16, mode="r")
    ix = torch.randint(len(data) - block_size - 1, (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, non_blocking=True), y.to(device, non_blocking=True)
    return x, y


# --------------------------------------------------------------------------
# Model: minimal nanoGPT-style decoder-only transformer
# --------------------------------------------------------------------------

class CausalSelfAttention(nn.Module):
    def __init__(self, n_embd, n_head, block_size, dropout):
        super().__init__()
        assert n_embd % n_head == 0
        self.n_head = n_head
        self.n_embd = n_embd
        self.c_attn = nn.Linear(n_embd, 3 * n_embd, bias=False)
        self.c_proj = nn.Linear(n_embd, n_embd, bias=False)
        self.attn_dropout = dropout
        self.resid_dropout = nn.Dropout(dropout)

    def forward(self, x):
        B, T, C = x.shape
        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        k = k.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)
        y = F.scaled_dot_product_attention(
            q, k, v, is_causal=True,
            dropout_p=self.attn_dropout if self.training else 0.0,
        )
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.resid_dropout(self.c_proj(y))


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

    def forward(self, x):
        return self.dropout(self.c_proj(self.gelu(self.c_fc(x))))


class Block(nn.Module):
    def __init__(self, n_embd, n_head, block_size, dropout):
        super().__init__()
        self.ln_1 = nn.LayerNorm(n_embd)
        self.attn = CausalSelfAttention(n_embd, n_head, block_size, dropout)
        self.ln_2 = nn.LayerNorm(n_embd)
        self.mlp = MLP(n_embd, dropout)

    def forward(self, x):
        x = x + self.attn(self.ln_1(x))
        x = x + self.mlp(self.ln_2(x))
        return x


class GPT(nn.Module):
    def __init__(self, vocab_size, block_size, n_layer, n_head, n_embd, dropout=0.0):
        super().__init__()
        self.block_size = block_size
        self.tok_emb = nn.Embedding(vocab_size, n_embd)
        self.pos_emb = nn.Embedding(block_size, n_embd)
        self.drop = nn.Dropout(dropout)
        self.blocks = nn.ModuleList(
            [Block(n_embd, n_head, block_size, dropout) for _ in range(n_layer)]
        )
        self.ln_f = nn.LayerNorm(n_embd)
        self.head = nn.Linear(n_embd, vocab_size, bias=False)
        self.tok_emb.weight = self.head.weight  # weight tying

        self.apply(self._init_weights)
        n_params = sum(p.numel() for p in self.parameters())
        print(f"Model: {n_params / 1e6:.1f}M parameters")

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

    def forward(self, idx, targets=None):
        B, T = idx.shape
        pos = torch.arange(T, device=idx.device)
        x = self.drop(self.tok_emb(idx) + self.pos_emb(pos))
        for block in self.blocks:
            x = block(x)
        x = self.ln_f(x)
        logits = self.head(x)
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
        return logits, loss

    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=0.8, top_k=50):
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -self.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)
            next_id = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, next_id), dim=1)
        return idx


# --------------------------------------------------------------------------
# Training
# --------------------------------------------------------------------------

def get_lr(it, warmup_iters, lr_decay_iters, max_lr, min_lr):
    if it < warmup_iters:
        return max_lr * (it + 1) / warmup_iters
    if it > lr_decay_iters:
        return min_lr
    ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters)
    coeff = 0.5 * (1.0 + math.cos(math.pi * ratio))
    return min_lr + coeff * (max_lr - min_lr)


def main():
    parser = argparse.ArgumentParser(description="Train a small GPT on the wiki text dump.")
    parser.add_argument("--data", type=str, default="wiki_dump.txt", help="Path to the text dump")
    parser.add_argument("--out-dir", type=str, default="out", help="Checkpoint output dir")
    parser.add_argument("--max-minutes", type=float, default=25, help="Hard time budget for training")
    parser.add_argument("--block-size", type=int, default=512, help="Context length")
    parser.add_argument("--batch-size", type=int, default=64, help="Batch size (fits 24GB at these dims)")
    parser.add_argument("--n-layer", type=int, default=12)
    parser.add_argument("--n-head", type=int, default=12)
    parser.add_argument("--n-embd", type=int, default=768)
    parser.add_argument("--lr", type=float, default=6e-4)
    parser.add_argument("--min-lr", type=float, default=6e-5)
    parser.add_argument("--weight-decay", type=float, default=0.1)
    parser.add_argument("--grad-clip", type=float, default=1.0)
    parser.add_argument("--eval-interval", type=int, default=200)
    parser.add_argument("--no-compile", action="store_true", help="Disable torch.compile")
    args = parser.parse_args()

    if not torch.cuda.is_available():
        print("No CUDA GPU found. This script needs a GPU (e.g. the rented 4090).")
        sys.exit(1)

    device = "cuda"
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True
    torch.manual_seed(1337)
    os.makedirs(args.out_dir, exist_ok=True)

    train_bin, val_bin = prepare_data(args.data)

    vocab_size = 50304  # round up from gpt2's 50257 to a multiple of 64 for faster matmuls
    model = GPT(
        vocab_size=vocab_size,
        block_size=args.block_size,
        n_layer=args.n_layer,
        n_head=args.n_head,
        n_embd=args.n_embd,
        dropout=0.0,
    ).to(device)

    if not args.no_compile:
        try:
            model = torch.compile(model)
            print("torch.compile enabled")
        except Exception as e:
            print(f"torch.compile unavailable ({e}), continuing without it")

    optimizer = torch.optim.AdamW(
        model.parameters(), lr=args.lr, weight_decay=args.weight_decay, betas=(0.9, 0.95)
    )

    # Rough iteration budget just for the LR schedule shape; the real stop
    # condition is the wall-clock timer below, not this count.
    warmup_iters = 100
    lr_decay_iters = 20000

    max_seconds = args.max_minutes * 60
    start_time = time.time()

    @torch.no_grad()
    def estimate_val_loss():
        model.eval()
        losses = []
        for _ in range(20):
            x, y = get_batch(val_bin, args.block_size, args.batch_size, device)
            with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
                _, loss = model(x, y)
            losses.append(loss.item())
        model.train()
        return sum(losses) / len(losses)

    model.train()
    it = 0
    pbar = tqdm(total=max_seconds, unit="s", desc="Training (time budget)")
    last_elapsed = 0.0

    while True:
        elapsed = time.time() - start_time
        if elapsed >= max_seconds:
            break

        lr = get_lr(it, warmup_iters, lr_decay_iters, args.lr, args.min_lr)
        for g in optimizer.param_groups:
            g["lr"] = lr

        x, y = get_batch(train_bin, args.block_size, args.batch_size, device)
        with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
            _, loss = model(x, y)

        optimizer.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
        optimizer.step()

        if it % args.eval_interval == 0:
            val_loss = estimate_val_loss()
            tqdm.write(f"iter {it}: train loss {loss.item():.4f}, val loss {val_loss:.4f}, lr {lr:.2e}")
            torch.save(
                {"model": model.state_dict(), "iter": it, "args": vars(args)},
                os.path.join(args.out_dir, "ckpt.pt"),
            )

        pbar.update(elapsed - last_elapsed)
        pbar.set_postfix(loss=f"{loss.item():.3f}", iter=it)
        last_elapsed = elapsed
        it += 1

    pbar.close()

    # Final checkpoint
    torch.save(
        {"model": model.state_dict(), "iter": it, "args": vars(args)},
        os.path.join(args.out_dir, "ckpt.pt"),
    )
    print(f"\nDone. Trained {it} iterations in {(time.time() - start_time) / 60:.1f} min. "
          f"Checkpoint saved to {os.path.join(args.out_dir, 'ckpt.pt')}")

    # Sample generation as a sanity check
    import tiktoken
    enc = tiktoken.get_encoding("gpt2")
    model.eval()
    prompt = "The history of"
    idx = torch.tensor([enc.encode_ordinary(prompt)], dtype=torch.long, device=device)
    out = model.generate(idx, max_new_tokens=150)
    print("\n--- Sample generation ---")
    print(enc.decode(out[0].tolist()))


if __name__ == "__main__":
    main()