File size: 4,330 Bytes
170658b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 | #!/usr/bin/env python3
"""Train a tiny character-level GPT on CPU.
This is intentionally small and educational, not production-grade.
"""
from __future__ import annotations
import argparse
import json
import time
from pathlib import Path
import torch
from model import TinyGPT, TinyGPTConfig
def build_vocab(text: str):
chars = sorted(set(text))
stoi = {ch: i for i, ch in enumerate(chars)}
itos = {i: ch for ch, i in stoi.items()}
return chars, stoi, itos
def encode(text: str, stoi: dict[str, int]):
return [stoi[ch] for ch in text]
def get_batch(data: torch.Tensor, block_size: int, batch_size: int, device: str):
ix = torch.randint(len(data) - block_size - 1, (batch_size,))
x = torch.stack([data[i : i + block_size] for i in ix]).to(device)
y = torch.stack([data[i + 1 : i + block_size + 1] for i in ix]).to(device)
return x, y
@torch.no_grad()
def estimate_loss(model, train_data, val_data, block_size, batch_size, eval_iters, device):
out = {}
model.eval()
for split, data in [("train", train_data), ("val", val_data)]:
losses = torch.zeros(eval_iters)
for k in range(eval_iters):
xb, yb = get_batch(data, block_size, batch_size, device)
_, loss = model(xb, yb)
losses[k] = loss.item()
out[split] = losses.mean().item()
model.train()
return out
def main():
p = argparse.ArgumentParser()
p.add_argument("--data", default="data/tiny_corpus.txt")
p.add_argument("--out", default="checkpoints/tinyllm.pt")
p.add_argument("--steps", type=int, default=500)
p.add_argument("--batch-size", type=int, default=16)
p.add_argument("--block-size", type=int, default=64)
p.add_argument("--n-layer", type=int, default=2)
p.add_argument("--n-head", type=int, default=2)
p.add_argument("--n-embd", type=int, default=64)
p.add_argument("--lr", type=float, default=3e-4)
p.add_argument("--eval-interval", type=int, default=100)
p.add_argument("--eval-iters", type=int, default=10)
p.add_argument("--seed", type=int, default=1337)
args = p.parse_args()
torch.manual_seed(args.seed)
device = "cpu"
data_path = Path(args.data)
text = data_path.read_text(encoding="utf-8")
if len(text) < args.block_size + 2:
raise SystemExit("Dataset is too small for the chosen block size.")
chars, stoi, itos = build_vocab(text)
encoded = torch.tensor(encode(text, stoi), dtype=torch.long)
n = int(0.9 * len(encoded))
train_data = encoded[:n]
val_data = encoded[n:] if len(encoded[n:]) > args.block_size + 1 else encoded[:n]
cfg = TinyGPTConfig(
vocab_size=len(chars),
block_size=args.block_size,
n_layer=args.n_layer,
n_head=args.n_head,
n_embd=args.n_embd,
dropout=0.1,
)
model = TinyGPT(cfg).to(device)
params = sum(p.numel() for p in model.parameters())
print(f"chars={len(chars)} tokens={len(encoded)} params={params:,} device={device}")
optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr)
start = time.time()
last_loss = None
for step in range(args.steps + 1):
if step % args.eval_interval == 0 or step == args.steps:
losses = estimate_loss(model, train_data, val_data, args.block_size, args.batch_size, args.eval_iters, device)
print(f"step {step:5d}: train {losses['train']:.4f}, val {losses['val']:.4f}")
last_loss = losses
xb, yb = get_batch(train_data, args.block_size, args.batch_size, device)
_, loss = model(xb, yb)
optimizer.zero_grad(set_to_none=True)
loss.backward()
optimizer.step()
out_path = Path(args.out)
out_path.parent.mkdir(parents=True, exist_ok=True)
ckpt = {
"model_state": model.state_dict(),
"config": cfg.__dict__,
"stoi": stoi,
"itos": {str(k): v for k, v in itos.items()},
"train_args": vars(args),
"last_loss": last_loss,
}
torch.save(ckpt, out_path)
meta_path = out_path.with_suffix(".json")
meta_path.write_text(json.dumps({"params": params, "chars": chars, "last_loss": last_loss}, indent=2), encoding="utf-8")
print(f"saved {out_path} and {meta_path} in {time.time() - start:.1f}s")
if __name__ == "__main__":
main()
|