#!/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()