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