train_compactlm5m.py

#!/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()