train.py

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