Add exact training script (defines CompactLM class)

train_compactlm5m.py

@@ -0,0 +1,328 @@
+#!/usr/bin/env python3
+"""
+CompactLM-5M — ~6.2M-param LLaMA-style English LM, from scratch.
+Requested by DedeProGames (model-requests #14): LLaMA-style, ~5M params,
+fineweb-edu, budget raised to ~100M tokens on the 30-50 min GPU window.
+
+Architecture (~6.16M params, tied embeddings):
+  - vocab 12288 (gollem_eval BPE, byte-level)
+  - d_model 256, n_layers 4, n_heads 4 (head_dim 64), SwiGLU ff 640
+  - RMSNorm pre-norm, RoPE, causal attention, ctx 512
+  - Standard LLaMA (no sliding window) — "LLaMA-style"
+
+Data: stream fineweb-edu (train) only. dclm-baseline-1.0 was failing
+(ConnectError) at build time on this host, so single-corpus — logged honestly.
+"""
+import os, sys, math, time, json, random, argparse, fcntl
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+TOK_PATH = "gollem_eval/tokenizer.json"
+CTX = 512
+
+def load_tok():
+    from tokenizers import Tokenizer
+    return Tokenizer.from_file(TOK_PATH)
+
+def stream_tokens(dsname, split, target_chars, tok, log):
+    from datasets import load_dataset
+    ds = load_dataset(dsname, split=split, streaming=True)
+    ids = []
+    nchars = 0
+    for row in ds:
+        text = row.get("text") or row.get("content") or ""
+        if not text:
+            continue
+        nchars += len(text)
+        ids.extend(tok.encode(text, add_special_tokens=False).ids)
+        if nchars >= target_chars:
+            break
+    return ids
+
+def precompute_rope(dim, max_pos, base=10000.0):
+    freqs = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
+    t = torch.arange(max_pos).float()
+    angles = torch.outer(t, freqs)
+    return torch.polar(torch.ones_like(angles), angles)
+
+def apply_rope(x, freqs_cis, offset=0):
+    B, nh, S, hd = x.shape
+    x = x.view(B, nh, S, hd // 2, 2)
+    xr = x[..., 0].float()
+    xi = x[..., 1].float()
+    fc = freqs_cis[offset:offset + S].to(x.device)
+    xr2 = xr * fc.real - xi * fc.imag
+    xi2 = xr * fc.imag + xi * fc.real
+    out = torch.stack([xr2, xi2], dim=-1).reshape(B, nh, S, hd)
+    return out.to(x.dtype)
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-5):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+    def forward(self, x):
+        norm = x.float().pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
+        return (x.float() * norm).to(x.dtype) * self.weight
+
+class Attention(nn.Module):
+    def __init__(self, d, n_heads):
+        super().__init__()
+        self.n_heads = n_heads
+        self.head_dim = d // n_heads
+        self.wq = nn.Linear(d, d, bias=False)
+        self.wk = nn.Linear(d, d, bias=False)
+        self.wv = nn.Linear(d, d, bias=False)
+        self.wo = nn.Linear(d, d, bias=False)
+    def forward(self, x, freqs_cis, offset=0):
+        B, S, _ = x.shape
+        q = self.wq(x).view(B, S, self.n_heads, self.head_dim).transpose(1, 2)
+        k = self.wk(x).view(B, S, self.n_heads, self.head_dim).transpose(1, 2)
+        v = self.wv(x).view(B, S, self.n_heads, self.head_dim).transpose(1, 2)
+        q = apply_rope(q, freqs_cis, offset)
+        k = apply_rope(k, freqs_cis, offset)
+        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        y = y.transpose(1, 2).reshape(B, S, -1)
+        return self.wo(y)
+
+class MLP(nn.Module):
+    def __init__(self, d, ff):
+        super().__init__()
+        self.w1 = nn.Linear(d, ff, bias=False)
+        self.w2 = nn.Linear(d, ff, bias=False)
+        self.w3 = nn.Linear(ff, d, bias=False)
+    def forward(self, x):
+        return self.w3(F.silu(self.w1(x)) * self.w2(x))
+
+class Block(nn.Module):
+    def __init__(self, d, n_heads, ff):
+        super().__init__()
+        self.ln1 = RMSNorm(d)
+        self.attn = Attention(d, n_heads)
+        self.ln2 = RMSNorm(d)
+        self.mlp = MLP(d, ff)
+    def forward(self, x, freqs_cis, offset=0):
+        x = x + self.attn(self.ln1(x), freqs_cis, offset)
+        x = x + self.mlp(self.ln2(x))
+        return x
+
+class CompactLM(nn.Module):
+    def __init__(self, vocab, d=256, n_layers=4, n_heads=4, ff=640, ctx=512):
+        super().__init__()
+        self.vocab = vocab; self.ctx = ctx; self.d = d
+        self.tok = nn.Embedding(vocab, d)
+        self.blocks = nn.ModuleList([Block(d, n_heads, ff) for _ in range(n_layers)])
+        self.ln_f = RMSNorm(d)
+        self.head = nn.Linear(d, vocab, bias=False)
+        self.head.weight = self.tok.weight  # tied
+        self.freqs_cis = precompute_rope(d // n_heads, ctx)
+        self.init_weights()
+    def init_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, mean=0.0, std=0.02)
+        nn.init.normal_(self.tok.weight, mean=0.0, std=0.02)
+    def forward(self, idx, targets=None):
+        B, S = idx.shape
+        h = self.tok(idx)
+        for b in self.blocks:
+            h = b(h, self.freqs_cis)
+        h = self.ln_f(h)
+        logits = self.head(h)
+        if targets is not None:
+            loss = F.cross_entropy(logits[:, :-1].reshape(-1, logits.size(-1)),
+                                   targets[:, 1:].reshape(-1), ignore_index=-1)
+            return loss
+        return logits
+    @torch.no_grad()
+    def generate(self, idx, max_new_tokens=128, temperature=0.8, top_k=40, seed=0):
+        g = torch.Generator(device=idx.device).manual_seed(seed)
+        for _ in range(max_new_tokens):
+            ctx_in = idx[:, -self.ctx:]
+            logits = self(ctx_in)[:, -1]
+            if temperature and temperature > 0:
+                logits = logits / temperature
+            if top_k:
+                v, _ = torch.topk(logits, top_k, dim=-1)
+                logits[logits < v[:, -1, None]] = float("-inf")
+            p = torch.softmax(logits, dim=1)
+            nxt = torch.multinomial(p, 1, generator=g)
+            idx = torch.cat([idx, nxt], dim=1)
+        return idx
+
+def count_params(m):
+    return sum(p.numel() for p in m.parameters())
+
+def build_data(args, tok, log):
+    outdir = os.path.join(args.out, "data")
+    os.makedirs(outdir, exist_ok=True)
+    train_npy = os.path.join(outdir, "train.npy")
+    val_npy = os.path.join(outdir, "val.npy")
+    if os.path.exists(train_npy) and os.path.exists(val_npy):
+        log(f"[data] reusing {train_npy}")
+        return
+    log("[data] streaming fineweb-edu ...")
+    t0 = time.time()
+    fw = stream_tokens("HuggingFaceFW/fineweb-edu", "train", args.fw_chars, tok, log)
+    log(f"[data] fineweb-edu: {len(fw):,} tokens from {args.fw_chars:,} chars in {time.time()-t0:.0f}s")
+    all_ids = np.array(fw, dtype=np.int32)
+    # NOTE: no token-level shuffle; batch_iter shuffles at the WINDOW level.
+    n_val = args.val_tokens
+    val = all_ids[:n_val]
+    train = all_ids[n_val:]
+    def pad(x):
+        n = (len(x) // CTX) * CTX
+        return x[:n].reshape(-1, CTX)
+    train = pad(train)
+    val = pad(val)
+    np.save(train_npy, train)
+    np.save(val_npy, val)
+    log(f"[data] train {train.shape} ({train.shape[0]*CTX:,} tok), val {val.shape} ({val.shape[0]*CTX:,} tok)")
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--out", default="models/compactlm-5m")
+    ap.add_argument("--steps", type=int, default=20000)
+    ap.add_argument("--batch", type=int, default=128)
+    ap.add_argument("--ctx", type=int, default=CTX)
+    ap.add_argument("--lr", type=float, default=3e-4)
+    ap.add_argument("--warmup", type=int, default=300)
+    ap.add_argument("--min-lr-frac", type=float, default=0.1)
+    ap.add_argument("--weight-decay", type=float, default=0.1)
+    ap.add_argument("--grad-clip", type=float, default=1.0)
+    ap.add_argument("--fw-chars", type=int, default=250_000_000)
+    ap.add_argument("--val-tokens", type=int, default=1_000_000)
+    ap.add_argument("--ckpt-every", type=int, default=500)
+    ap.add_argument("--val-every", type=int, default=1000)
+    ap.add_argument("--sample-every", type=int, default=2000)
+    ap.add_argument("--resume", action="store_true")
+    args = ap.parse_args()
+
+    os.makedirs(args.out, exist_ok=True)
+    logf = open(os.path.join(args.out, "train.log"), "a")
+    def log(s):
+        print(s, flush=True)
+        logf.write(s + "\n"); logf.flush()
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    log(f"[init] device={device}")
+    if device == "cuda":
+        free, total = torch.cuda.mem_get_info()
+        log(f"[init] gpu free {free/1e6:.1f} / {total/1e6:.1f} MB")
+
+    tok = load_tok()
+    vocab = tok.get_vocab_size()
+    log(f"[init] vocab={vocab}")
+
+    build_data(args, tok, log)
+
+    model = CompactLM(vocab, d=256, n_layers=4, n_heads=4, ff=640, ctx=args.ctx).to(device)
+    npar = count_params(model)
+    log(f"[init] params={npar:,} ({npar/1e6:.2f}M)")
+    assert 5_500_000 <= npar <= 7_000_000, f"param count {npar} not ~6M"
+
+    train_ids = np.load(os.path.join(args.out, "data", "train.npy"))
+    val_ids = np.load(os.path.join(args.out, "data", "val.npy"))
+
+    def batch_iter(ids, batch, shuffle=True, seed=0):
+        n = ids.shape[0]
+        rng = np.random.default_rng(seed)
+        idx = rng.permutation(n) if shuffle else np.arange(n)
+        for start in range(0, n, batch):
+            sel = idx[start:start + batch]
+            if len(sel) < batch:
+                continue
+            yield torch.from_numpy(ids[sel]).long().to(device)
+
+    step = 0
+    best_val = float("inf")
+    if args.resume and os.path.exists(os.path.join(args.out, "last.pt")):
+        ck = torch.load(os.path.join(args.out, "last.pt"), map_location="cpu")
+        model.load_state_dict(ck["model"])
+        step = ck["step"]; best_val = ck.get("best_val", float("inf"))
+        log(f"[resume] from step {step}, best_val {best_val:.4f}")
+
+    no_decay = [p for n, p in model.named_parameters() if p.ndim <= 1]
+    decay = [p for n, p in model.named_parameters() if p.ndim > 1]
+    groups = [{"params": decay, "weight_decay": args.weight_decay},
+              {"params": no_decay, "weight_decay": 0.0}]
+    opt = torch.optim.AdamW(groups, lr=args.lr, betas=(0.9, 0.95), eps=1e-8)
+
+    def lr_at(step):
+        if step < args.warmup:
+            return args.lr * (step + 1) / args.warmup
+        p = (step - args.warmup) / max(1, args.steps - args.warmup)
+        p = min(1.0, p)
+        return args.lr * (args.min_lr_frac + (1 - args.min_lr_frac) * 0.5 * (1 + math.cos(math.pi * p)))
+
+    model.train()
+    t0 = time.time()
+    train_iter = None
+    while step < args.steps:
+        if train_iter is None:
+            train_iter = batch_iter(train_ids, args.batch, shuffle=True, seed=step // 1000)
+        try:
+            b = next(train_iter)
+        except StopIteration:
+            train_iter = batch_iter(train_ids, args.batch, shuffle=True, seed=step // 1000)
+            b = next(train_iter)
+        for g in groups:
+            g["lr"] = lr_at(step)
+        opt.zero_grad()
+        loss = model(b, b)
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+        opt.step()
+        step += 1
+        if step % 100 == 0 or step == 1:
+            tok_s = (step * args.batch * args.ctx) / max(1e-6, time.time() - t0)
+            log(f"[step {step}/{args.steps}] loss {loss.item():.4f} lr {lr_at(step):.2e} tok/s {tok_s:,.0f}")
+        if step % args.ckpt_every == 0:
+            torch.save({"model": model.state_dict(), "step": step, "best_val": best_val, "vocab": vocab},
+                       os.path.join(args.out, "last.pt"))
+        if step % args.val_every == 0:
+            model.eval()
+            with torch.no_grad():
+                vloss = 0.0; n = 0
+                for b in batch_iter(val_ids, 32, shuffle=False, seed=0):
+                    vloss += model(b, b).item(); n += 1
+                vloss /= max(1, n)
+            log(f"[val step {step}] val_loss {vloss:.4f} ppl {math.exp(min(vloss,20)):.2f}")
+            if vloss < best_val:
+                best_val = vloss
+                torch.save({"model": model.state_dict(), "step": step, "best_val": best_val, "vocab": vocab},
+                           os.path.join(args.out, "best.pt"))
+                log(f"[val step {step}] NEW BEST -> best.pt")
+            model.train()
+        if step % args.sample_every == 0:
+            model.eval()
+            with torch.no_grad():
+                tok_prompts = ["Once upon a time", "The cat sat on the", "def hello():"]
+                for p in tok_prompts:
+                    ids = torch.tensor([tok.encode(p, add_special_tokens=False).ids], device=device)
+                    out = model.generate(ids, max_new_tokens=80, temperature=0.8, top_k=40, seed=step % 1000)
+                    log(f"[sample step {step}] {tok.decode(out[0].tolist(), skip_special_tokens=True)[:400]!r}")
+            model.train()
+    torch.save({"model": model.state_dict(), "step": step, "best_val": best_val, "vocab": vocab},
+               os.path.join(args.out, "final.pt"))
+    model.eval()
+    with torch.no_grad():
+        vloss = 0.0; n = 0
+        for b in batch_iter(val_ids, 32, shuffle=False, seed=0):
+            vloss += model(b, b).item(); n += 1
+        vloss /= max(1, n)
+    log(f"[final] step {step} val_loss {vloss:.4f} ppl {math.exp(min(vloss,20)):.2f}")
+    logf.close()
+
+if __name__ == "__main__":
+    _inherited_fd = os.environ.get("CLM_LOCK_FD")
+    if _inherited_fd is None:
+        _lock = open(os.path.join(os.path.dirname(os.path.abspath(__file__)), ".train_clm5m.lock"), "w")
+        try:
+            fcntl.flock(_lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
+        except BlockingIOError:
+            print("[lock] another compactlm5m process already holds the lock — REFUSING to double-launch", flush=True)
+            sys.exit(3)
+    main()