1 files changed, 183 insertions, 0 deletions
diff --git a/ep_run/train.py b/ep_run/train.py
new file mode 100644
index 0000000..90c5361
--- /dev/null
+++ b/ep_run/train.py
@@ -0,0 +1,183 @@
+"""Train a tiny char-level GPT on tinyshakespeare with softmax or sigmoid attn.
+
+Logs train/val loss to runs/<run_name>/log.jsonl (one JSON per line).
+No checkpoints are saved (disk-conscious).
+"""
+import argparse
+import json
+import math
+import os
+import pickle
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from model import GPT, GPTConfig
+
+
+def get_batch(split: str, data_dir: Path, block_size: int, batch_size: int, device: str):
+    fn = "train.bin" if split == "train" else "val.bin"
+    data = np.memmap(data_dir / fn, dtype=np.uint16, mode="r")
+    ix = torch.randint(len(data) - block_size - 1, (batch_size,))
+    x = torch.stack([torch.from_numpy(data[i : i + block_size].astype(np.int64)) for i in ix])
+    y = torch.stack([torch.from_numpy(data[i + 1 : i + 1 + block_size].astype(np.int64)) for i in ix])
+    return x.to(device, non_blocking=True), y.to(device, non_blocking=True)
+
+
+@torch.no_grad()
+def estimate_loss(model, data_dir, block_size, batch_size, device, eval_iters):
+    out = {}
+    model.eval()
+    for split in ("train", "val"):
+        losses = torch.zeros(eval_iters)
+        for k in range(eval_iters):
+            X, Y = get_batch(split, data_dir, block_size, batch_size, device)
+            _, loss = model(X, Y)
+            losses[k] = loss.item()
+        out[split] = losses.mean().item()
+    model.train()
+    return out
+
+
+def lr_schedule(it, warmup_iters, lr_decay_iters, max_lr, min_lr):
+    if it < warmup_iters:
+        return max_lr * (it + 1) / (warmup_iters + 1)
+    if it > lr_decay_iters:
+        return min_lr
+    decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters)
+    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
+    return min_lr + coeff * (max_lr - min_lr)
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--run_name", type=str, required=True)
+    p.add_argument("--attn_mode", type=str, default="softmax", choices=["softmax", "sigmoid"])
+    p.add_argument("--sigmoid_bias_mode", type=str, default="neg_log_n",
+                   choices=["zero", "neg_log_n", "learned"])
+    p.add_argument("--seed", type=int, default=1337)
+    p.add_argument("--data_dir", type=str, default="data/shakespeare_char")
+    p.add_argument("--out_dir", type=str, default="runs")
+    p.add_argument("--block_size", type=int, default=256)
+    p.add_argument("--batch_size", type=int, default=64)
+    p.add_argument("--n_layer", type=int, default=6)
+    p.add_argument("--n_head", type=int, default=6)
+    p.add_argument("--n_embd", type=int, default=384)
+    p.add_argument("--dropout", type=float, default=0.2)
+    p.add_argument("--max_iters", type=int, default=5000)
+    p.add_argument("--warmup_iters", type=int, default=100)
+    p.add_argument("--lr_decay_iters", type=int, default=5000)
+    p.add_argument("--max_lr", type=float, default=1e-3)
+    p.add_argument("--min_lr", type=float, default=1e-4)
+    p.add_argument("--weight_decay", type=float, default=0.1)
+    p.add_argument("--beta1", type=float, default=0.9)
+    p.add_argument("--beta2", type=float, default=0.99)
+    p.add_argument("--grad_clip", type=float, default=1.0)
+    p.add_argument("--eval_interval", type=int, default=250)
+    p.add_argument("--eval_iters", type=int, default=100)
+    p.add_argument("--log_interval", type=int, default=50)
+    p.add_argument("--dtype", type=str, default="bfloat16", choices=["float32", "bfloat16", "float16"])
+    args = p.parse_args()
+
+    torch.manual_seed(args.seed)
+    torch.cuda.manual_seed_all(args.seed)
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    data_dir = Path(args.data_dir)
+    with open(data_dir / "meta.pkl", "rb") as f:
+        meta = pickle.load(f)
+    vocab_size = meta["vocab_size"]
+
+    run_dir = Path(args.out_dir) / args.run_name
+    run_dir.mkdir(parents=True, exist_ok=True)
+    log_path = run_dir / "log.jsonl"
+    cfg_path = run_dir / "config.json"
+    with open(cfg_path, "w") as f:
+        json.dump(vars(args) | {"vocab_size": vocab_size}, f, indent=2)
+
+    cfg = GPTConfig(
+        block_size=args.block_size,
+        vocab_size=vocab_size,
+        n_layer=args.n_layer,
+        n_head=args.n_head,
+        n_embd=args.n_embd,
+        dropout=args.dropout,
+        attn_mode=args.attn_mode,
+        sigmoid_bias_mode=args.sigmoid_bias_mode,
+    )
+    model = GPT(cfg).to(device)
+    n_params = model.num_params()
+
+    # AdamW with weight-decay-free for 1D params (ln, embeddings, biases, sig_bias)
+    decay_params, nodecay_params = [], []
+    for n, pr in model.named_parameters():
+        if not pr.requires_grad:
+            continue
+        if pr.dim() >= 2:
+            decay_params.append(pr)
+        else:
+            nodecay_params.append(pr)
+    optimizer = torch.optim.AdamW(
+        [
+            {"params": decay_params, "weight_decay": args.weight_decay},
+            {"params": nodecay_params, "weight_decay": 0.0},
+        ],
+        lr=args.max_lr,
+        betas=(args.beta1, args.beta2),
+        fused=(device == "cuda"),
+    )
+
+    dtype_map = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16}
+    amp_dtype = dtype_map[args.dtype]
+    scaler = torch.amp.GradScaler("cuda", enabled=(args.dtype == "float16"))
+
+    t0 = time.time()
+
+    def log(record: dict):
+        record["t"] = time.time() - t0
+        with open(log_path, "a") as f:
+            f.write(json.dumps(record) + "\n")
+
+    log({"event": "start", "params": n_params, "config": vars(args) | {"vocab_size": vocab_size}})
+    print(f"[{args.run_name}] params={n_params/1e6:.2f}M device={device} dtype={args.dtype}")
+
+    model.train()
+    for it in range(args.max_iters + 1):
+        lr = lr_schedule(it, args.warmup_iters, args.lr_decay_iters, args.max_lr, args.min_lr)
+        for g in optimizer.param_groups:
+            g["lr"] = lr
+
+        if it % args.eval_interval == 0 or it == args.max_iters:
+            losses = estimate_loss(model, data_dir, args.block_size, args.batch_size, device, args.eval_iters)
+            log({"event": "eval", "iter": it, "train_loss": losses["train"], "val_loss": losses["val"], "lr": lr})
+            print(f"[{args.run_name}] iter {it:5d}  train {losses['train']:.4f}  val {losses['val']:.4f}  lr {lr:.4g}")
+
+        if it == args.max_iters:
+            break
+
+        X, Y = get_batch("train", data_dir, args.block_size, args.batch_size, device)
+        with torch.amp.autocast(device_type="cuda", dtype=amp_dtype, enabled=(device == "cuda")):
+            _, loss = model(X, Y)
+        optimizer.zero_grad(set_to_none=True)
+        if args.dtype == "float16":
+            scaler.scale(loss).backward()
+            scaler.unscale_(optimizer)
+            torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+            scaler.step(optimizer)
+            scaler.update()
+        else:
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+            optimizer.step()
+
+        if it % args.log_interval == 0:
+            log({"event": "step", "iter": it, "train_loss": loss.item(), "lr": lr})
+
+    log({"event": "done", "iter": args.max_iters})
+    print(f"[{args.run_name}] done in {time.time()-t0:.1f}s")
+
+
+if __name__ == "__main__":
+    main()