import argparse import math import os from pathlib import Path import torch import torch.nn as nn from torch.nn import functional as F DEFAULT_CHARS = ( "\n" " " "abcdefghijklmnopqrstuvwxyz" "ABCDEFGHIJKLMNOPQRSTUVWXYZ" "0123456789" ".,!?;:'\"-_/\\()[]{}<>@#$%^&*+=|`~" ) class TinyTransformerLM(nn.Module): def __init__(self, vocab_size, block_size, n_embd=128, n_head=2, n_layer=2, dropout=0.1): super().__init__() self.block_size = block_size self.token_embedding = nn.Embedding(vocab_size, n_embd) self.position_embedding = nn.Embedding(block_size, n_embd) encoder_layer = nn.TransformerEncoderLayer( d_model=n_embd, nhead=n_head, dim_feedforward=4 * n_embd, dropout=dropout, activation="gelu", batch_first=True, ) self.blocks = nn.TransformerEncoder(encoder_layer, num_layers=n_layer) self.ln_f = nn.LayerNorm(n_embd) self.head = nn.Linear(n_embd, vocab_size) def forward(self, idx, targets=None): batch, time = idx.shape if time > self.block_size: raise ValueError("sequence is longer than block_size") token_emb = self.token_embedding(idx) pos = torch.arange(time, device=idx.device) pos_emb = self.position_embedding(pos)[None, :, :] x = token_emb + pos_emb mask = torch.triu(torch.ones(time, time, device=idx.device), diagonal=1).bool() x = self.blocks(x, mask=mask) x = self.ln_f(x) logits = self.head(x) loss = None if targets is not None: loss = F.cross_entropy(logits.reshape(batch * time, -1), targets.reshape(batch * time)) return logits, loss PRESETS = { "tiny": {"block_size": 64, "n_embd": 64, "n_head": 2, "n_layer": 1, "batch_size": 4, "steps": 1200, "lr": 3e-4}, "turbo": {"block_size": 32, "n_embd": 64, "n_head": 4, "n_layer": 2, "batch_size": 16, "steps": 600, "lr": 1e-3}, "fast": {"block_size": 64, "n_embd": 96, "n_head": 3, "n_layer": 2, "batch_size": 8, "steps": 800, "lr": 5e-4}, "smart": {"block_size": 128, "n_embd": 160, "n_head": 4, "n_layer": 3, "batch_size": 12, "steps": 1500, "lr": 3e-4}, "power": {"block_size": 128, "n_embd": 256, "n_head": 8, "n_layer": 4, "batch_size": 16, "steps": 1000, "lr": 4e-4}, "small": {"block_size": 128, "n_embd": 128, "n_head": 2, "n_layer": 2, "batch_size": 8, "steps": 1200, "lr": 3e-4}, "big": {"block_size": 128, "n_embd": 192, "n_head": 4, "n_layer": 4, "batch_size": 4, "steps": 1200, "lr": 2e-4}, "large": {"block_size": 128, "n_embd": 256, "n_head": 8, "n_layer": 6, "batch_size": 2, "steps": 1200, "lr": 1.5e-4}, } def build_vocab(text): chars = sorted(set(DEFAULT_CHARS + text)) stoi = {ch: i for i, ch in enumerate(chars)} itos = {i: ch for ch, i in stoi.items()} return stoi, itos def encode_text(text, stoi): fallback = stoi.get(" ", 0) return torch.tensor([stoi.get(ch, fallback) for ch in text], dtype=torch.long) def make_batch(data, batch_size, block_size, device): max_start = len(data) - block_size - 1 starts = torch.randint(max_start, (batch_size,)) x = torch.stack([data[i : i + block_size] for i in starts]) y = torch.stack([data[i + 1 : i + block_size + 1] for i in starts]) return x.to(device), y.to(device) @torch.no_grad() def estimate_loss(model, train_data, val_data, batch_size, block_size, device, eval_iters=20): model.eval() out = {} for split, data in (("train", train_data), ("val", val_data)): losses = [] for _ in range(eval_iters): x, y = make_batch(data, batch_size, block_size, device) _, loss = model(x, y) losses.append(loss.item()) out[split] = sum(losses) / len(losses) model.train() return out def main(): parser = argparse.ArgumentParser() parser.add_argument("--data", default="data/input.txt") parser.add_argument("--out", default="runs/tiny-char-model.pt") parser.add_argument("--preset", choices=sorted(PRESETS), default="tiny") parser.add_argument("--steps", type=int, default=1200) parser.add_argument("--batch-size", type=int, default=16) parser.add_argument("--block-size", type=int, default=128) parser.add_argument("--n-embd", type=int, default=128) parser.add_argument("--n-head", type=int, default=2) parser.add_argument("--n-layer", type=int, default=2) parser.add_argument("--lr", type=float, default=3e-4) args = parser.parse_args() preset = PRESETS[args.preset] if args.steps == 1200: args.steps = preset["steps"] if args.batch_size == 16: args.batch_size = preset["batch_size"] if args.block_size == 128: args.block_size = preset["block_size"] if args.n_embd == 128: args.n_embd = preset["n_embd"] if args.n_head == 2: args.n_head = preset["n_head"] if args.n_layer == 2: args.n_layer = preset["n_layer"] if args.lr == 3e-4: args.lr = preset["lr"] device = "cuda" if torch.cuda.is_available() else "cpu" if device == "cpu": threads = os.cpu_count() or 4 # Use the actual number of logical CPU cores torch.set_num_threads(threads) torch.set_num_interop_threads(1) torch.set_float32_matmul_precision("high") print(f"CPU optimization: using {threads} threads") else: torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.benchmark = True torch.set_float32_matmul_precision("high") print("CUDA optimization: TF32 enabled, cudnn benchmark on") text = Path(args.data).read_text(encoding="utf-8") stoi, itos = build_vocab(text) encoded = encode_text(text, stoi) if len(encoded) < args.block_size + 2: raise SystemExit("Dataset is too small. Add more text or lower --block-size.") split = max(1, int(0.9 * len(encoded))) train_data = encoded[:split] val_data = encoded[split - args.block_size - 1 :] chars = [ch for ch, _ in sorted(stoi.items(), key=lambda item: item[1])] model = TinyTransformerLM( vocab_size=len(chars), block_size=args.block_size, n_embd=args.n_embd, n_head=args.n_head, n_layer=args.n_layer, ).to(device) optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr) scaler = torch.cuda.amp.GradScaler(enabled=device == "cuda") params = sum(p.numel() for p in model.parameters()) print(f"device={device} params={params:,} vocab={len(chars)}") for step in range(args.steps + 1): if step % 100 == 0: losses = estimate_loss(model, train_data, val_data, args.batch_size, args.block_size, device) ppl = math.exp(min(losses["val"], 20)) print(f"step {step:5d} train {losses['train']:.4f} val {losses['val']:.4f} ppl {ppl:.2f}") xb, yb = make_batch(train_data, args.batch_size, args.block_size, device) if device == "cuda": with torch.cuda.amp.autocast(): _, loss = model(xb, yb) else: _, loss = model(xb, yb) optimizer.zero_grad(set_to_none=True) scaler.scale(loss).backward() if device == "cuda" else loss.backward() if device == "cuda": scaler.step(optimizer) scaler.update() else: optimizer.step() out_path = Path(args.out) out_path.parent.mkdir(parents=True, exist_ok=True) torch.save( { "model": model.state_dict(), "config": { "vocab_size": len(chars), "block_size": args.block_size, "n_embd": args.n_embd, "n_head": args.n_head, "n_layer": args.n_layer, }, "stoi": stoi, "itos": itos, }, out_path, ) print(f"saved {out_path}") if __name__ == "__main__": main()