"""Train one nano-GPT per (corpus, scheme) cell. python src/train.py # train everything that's missing python src/train.py --force # retrain all python src/train.py --quick # fast smoke run (few iters) python src/train.py --only-corpus shakespeare --only-scheme bpe512 For each cell it: builds + saves the tokenizer, encodes the (capped) corpus, trains a GPT, prints a sample, and saves model.pt + config.json into artifacts///. """ from __future__ import annotations import argparse import json import math import os import time import torch import config from model import GPT, GPTConfig from tokenizer import build_tokenizer, load_tokenizer # --------------------------------------------------------------------------- # def get_batch(data, block_size, batch_size, device): ix = torch.randint(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, non_blocking=True), y.to(device, non_blocking=True) @torch.no_grad() def estimate_loss(model, splits, args, device): out = {} model.eval() for name, data in splits.items(): losses = torch.zeros(args.eval_iters) for k in range(args.eval_iters): x, y = get_batch(data, args.block_size, args.batch_size, device) with _autocast(device): _, loss = model(x, y) losses[k] = loss.item() out[name] = losses.mean().item() model.train() return out def _autocast(device): if device.type == "cuda": return torch.autocast(device_type="cuda", dtype=torch.bfloat16) return torch.autocast(device_type="cpu", enabled=False) def lr_at(it, args): if it < args.warmup: return args.lr * (it + 1) / args.warmup if it > args.iters: return args.min_lr ratio = (it - args.warmup) / max(1, args.iters - args.warmup) coeff = 0.5 * (1.0 + math.cos(math.pi * ratio)) return args.min_lr + coeff * (args.lr - args.min_lr) def cell_is_done(adir: str) -> bool: return all(os.path.exists(os.path.join(adir, f)) for f in ("tokenizer.json", "model.pt", "config.json")) # --------------------------------------------------------------------------- # def train_cell(corpus: dict, scheme: dict, args, device) -> None: adir = config.artifact_dir(corpus["name"], scheme["name"]) tag = f"{corpus['name']}/{scheme['name']}" if cell_is_done(adir) and not args.force: print(f"[{tag}] already trained — skip (use --force to retrain)") return text_path = config.corpus_text_path(corpus["name"]) if not os.path.exists(text_path): print(f"[{tag}] corpus text missing ({text_path}); run prepare_data.py first") return text = open(text_path, encoding="utf-8", errors="replace").read() if args.max_chars and len(text) > args.max_chars: text = text[:args.max_chars] print(f"\n[{tag}] building tokenizer over {len(text)/1e6:.2f}M chars ...") t0 = time.time() tok = build_tokenizer(scheme, text, verbose=args.verbose) os.makedirs(adir, exist_ok=True) tok.save(adir) print(f"[{tag}] tokenizer: vocab={tok.vocab_size} " f"merges={len(getattr(tok, 'merges', []))} ({time.time()-t0:.1f}s)") # encode corpus -> ids t0 = time.time() ids = torch.tensor(tok.encode(text), dtype=torch.long) n = len(ids) tok_per_char = n / len(text) split = int(0.9 * n) splits = {"train": ids[:split], "val": ids[split:]} print(f"[{tag}] encoded {n} tokens " f"({tok_per_char:.2f} tok/char) ({time.time()-t0:.1f}s)") # model block_size = min(args.block_size, len(splits["val"]) - 1) cfg = GPTConfig(vocab_size=tok.vocab_size, block_size=block_size, n_layer=args.n_layer, n_head=args.n_head, n_embd=args.n_embd, dropout=args.dropout) model = GPT(cfg).to(device) print(f"[{tag}] model params: {model.num_params()/1e6:.2f}M block={block_size}") opt = torch.optim.AdamW(model.parameters(), lr=args.lr, betas=(0.9, 0.99), weight_decay=0.1) t0 = time.time() model.train() for it in range(args.iters + 1): for g in opt.param_groups: g["lr"] = lr_at(it, args) if it % args.eval_interval == 0 or it == args.iters: losses = estimate_loss(model, splits, args, device) print(f"[{tag}] iter {it:5d}/{args.iters} " f"train {losses['train']:.3f} val {losses['val']:.3f} " f"lr {opt.param_groups[0]['lr']:.2e} ({time.time()-t0:.0f}s)") if it == args.iters: break x, y = get_batch(splits["train"], block_size, args.batch_size, device) with _autocast(device): _, loss = model(x, y) opt.zero_grad(set_to_none=True) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) opt.step() # sample sample = generate_sample(model, tok, corpus.get("sample_text", "\n"), device) print(f"[{tag}] sample: {sample!r}") # bits-per-character makes loss comparable across tokenizers (per-token loss # is not — a bigger vocab inflates it). lower is better, and BPE usually wins. bits_per_char = losses["val"] * tok_per_char / math.log(2) # save torch.save(model.state_dict(), os.path.join(adir, "model.pt")) with open(os.path.join(adir, "config.json"), "w", encoding="utf-8") as f: json.dump({ "corpus": corpus["name"], "scheme": scheme["name"], "model_config": cfg.to_dict(), "tokenizer": tok.to_meta(), "final_loss": losses, "tok_per_char": tok_per_char, "bits_per_char": bits_per_char, "sample": sample, }, f, ensure_ascii=False, indent=2) print(f"[{tag}] saved -> {adir} (val {losses['val']:.3f} | {bits_per_char:.2f} bits/char)") @torch.no_grad() def generate_sample(model, tok, seed_text, device, n_new=200): model.eval() ids = tok.encode(seed_text) or tok.encode("\n") or [1] idx = torch.tensor([ids], dtype=torch.long, device=device) out = model.generate(idx, max_new_tokens=n_new, temperature=0.8, top_k=40) model.train() return tok.decode(out[0].tolist()) # --------------------------------------------------------------------------- # def build_args(): ap = argparse.ArgumentParser() ap.add_argument("--force", action="store_true") ap.add_argument("--quick", action="store_true", help="few iters, for smoke testing") ap.add_argument("--verbose", action="store_true") ap.add_argument("--only-corpus", default=None) ap.add_argument("--only-scheme", default=None) ap.add_argument("--device", default=None) ap.add_argument("--max-chars", type=int, default=1_500_000) ap.add_argument("--iters", type=int, default=3000) ap.add_argument("--batch-size", type=int, default=64) ap.add_argument("--block-size", type=int, default=128) ap.add_argument("--n-layer", type=int, default=4) ap.add_argument("--n-head", type=int, default=4) ap.add_argument("--n-embd", type=int, default=128) ap.add_argument("--dropout", type=float, default=0.1) ap.add_argument("--lr", type=float, default=3e-4) ap.add_argument("--min-lr", type=float, default=3e-5) ap.add_argument("--warmup", type=int, default=100) ap.add_argument("--eval-interval", type=int, default=250) ap.add_argument("--eval-iters", type=int, default=50) args = ap.parse_args() if args.quick: args.iters = 200 args.eval_interval = 100 args.warmup = 20 args.max_chars = min(args.max_chars, 400_000) return args def main(): args = build_args() device = torch.device(args.device or ("cuda" if torch.cuda.is_available() else "cpu")) torch.manual_seed(1337) print(f"device: {device} | iters: {args.iters} | max_chars: {args.max_chars}") corpora = [c for c in config.corpora() if not args.only_corpus or c["name"] == args.only_corpus] schemes = [s for s in config.schemes() if not args.only_scheme or s["name"] == args.only_scheme] t0 = time.time() for corpus in corpora: for scheme in schemes: train_cell(corpus, scheme, args, device) print(f"\nAll cells done in {time.time()-t0:.0f}s.") if __name__ == "__main__": main()