Upload train.py with huggingface_hub

train.py

@@ -0,0 +1,216 @@
+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()