|
|
| """
|
| 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-----"
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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()
|
|
|
|
|
| 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)
|
|
|
| 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
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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
|
|
|
| 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
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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
|
| 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)
|
| )
|
|
|
|
|
|
|
| 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()
|
|
|
|
|
| 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')}")
|
|
|
|
|
| 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() |
