| 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 |
| 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() |
|
|
