"""Causal LM pretraining for TinyLiquid on packed uint16 token files. Usage: .venv/bin/python train/train_lm.py --data data/train.bin --val data/valid.bin \ --config tiny10m --ckpt ckpt/nlp --steps 12000 Resume: .venv/bin/python train/train_lm.py ... --resume ckpt/nlp """ import argparse import json import math import os import random import sys import time from pathlib import Path import numpy as np import torch import torch.nn.functional as F from model.config import TinyLiquidConfig, CONFIGS from model.utils import latest_ckpt from model.tiny_liquid import TinyLiquid from data.tokenizer import load_tokenizer, PERSONA_TOKENS SAVE_PREFIX = "model" def parse_args(): ap = argparse.ArgumentParser() ap.add_argument("--data", default="data/train.bin") ap.add_argument("--val", default="data/valid.bin") ap.add_argument("--tok", default="data/tokenizer.json") ap.add_argument("--config", default="tiny10m") ap.add_argument("--ckpt", default="ckpt/nlp") ap.add_argument("--resume", default=None) ap.add_argument("--init-from", default=None, help="load model weights only (fresh optimizer + LR schedule)") ap.add_argument("--batch", type=int, default=16) ap.add_argument("--seq", type=int, default=256) ap.add_argument("--lr", type=float, default=3e-4) ap.add_argument("--min-lr", type=float, default=1e-5) ap.add_argument("--warmup", type=int, default=200) ap.add_argument("--steps", type=int, default=12000) ap.add_argument("--total-steps", type=int, default=None, help="absolute final step when resuming; overrides additive --steps") ap.add_argument("--log-every", type=int, default=25) ap.add_argument("--eval-every", type=int, default=500) ap.add_argument("--save-every", type=int, default=1000) ap.add_argument("--val-batches", type=int, default=40) ap.add_argument("--sample-max-new", type=int, default=80) ap.add_argument("--seed", type=int, default=42) ap.add_argument("--threads", type=int, default=8) ap.add_argument("--bf16", action="store_true", help="CPU autocast bf16 (SVE2/BF16 ARMv9)") ap.add_argument("--nan-rollback", type=int, default=50, help="auto-resume from last checkpoint after N consecutive non-finite steps (0=off)") ap.add_argument("--mtp", type=int, default=0, help="multi-token prediction aux heads (Meta MTP; 0=off)") return ap.parse_args() def load_bin(path, memmap_threshold=300 * 1024 * 1024): size = Path(path).stat().st_size if size >= memmap_threshold: # memory-map large corpora; sample_batch converts slices on the fly return np.memmap(path, dtype=np.uint16, mode="r") arr = np.fromfile(path, dtype=np.uint16) return torch.from_numpy(arr.astype(np.int64)) def make_batches(tokens: torch.Tensor, batch: int, seq: int, rng: random.Random): # non-overlapping windows; cycle through them in order n = (len(tokens) - 1) // seq if n <= 0: raise ValueError("corpus too small for seq") return torch.arange(0, n * seq, seq, dtype=torch.long) def sample_batch(tokens: torch.Tensor, offsets, step: int, batch: int, seq: int): pos = (step * batch) % len(offsets) idx = offsets[pos : pos + batch] if len(idx) < batch: idx = torch.cat([idx, offsets[: batch - len(idx)]]) if isinstance(tokens, np.memmap): buf = torch.stack([torch.from_numpy(tokens[int(s): int(s) + seq].astype(np.int64)) for s in idx]) else: buf = torch.stack([tokens[int(s): int(s) + seq] for s in idx]) return buf[:, :-1], buf[:, 1:] @torch.no_grad() def evaluate(model, val_tokens, batch, seq, n_batches, rng, persona=0, bf16=False): model.eval() total, count = 0.0, 0 offsets = make_batches(val_tokens, batch, seq, rng) for i in range(n_batches): x, y = sample_batch(val_tokens, offsets, i, batch, seq) p = torch.full((batch,), persona, dtype=torch.long) if persona else None if bf16: with torch.autocast("cpu", dtype=torch.bfloat16): logits = model(x, persona_ids=p) loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), y.reshape(-1)) else: logits = model(x, persona_ids=p) loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), y.reshape(-1)) total += loss.item() * y.numel() count += y.numel() model.train() return total / count @torch.no_grad() def sample(model, tok, prompt, max_new, persona=0, temp=0.8, top_k=40): ids = tok.encode(prompt).ids out = model.generate(tok, ids, persona_id=persona, max_new=max_new, temperature=temp, top_k=top_k) return tok.decode(out[len(ids):]) def main(): args = parse_args() torch.set_num_threads(args.threads) torch.manual_seed(args.seed) random.seed(args.seed) rng = random.Random(args.seed) ckpt_dir = Path(args.ckpt) ckpt_dir.mkdir(parents=True, exist_ok=True) tok = load_tokenizer(args.tok) cfg = TinyLiquidConfig(vocab_size=tok.get_vocab_size(), mtp_heads=args.mtp, **CONFIGS[args.config]) model = TinyLiquid(cfg) opt = torch.optim.AdamW(model.parameters(), lr=args.lr, betas=(0.9, 0.95), weight_decay=0.1) step, best_val = 0, float("inf") grad_accum = 1 # keep simple on CPU resume_from = 0 if args.init_from: path = latest_ckpt(args.init_from) if path: sd = torch.load(path, map_location="cpu") if args.mtp and set(sd["model"]) != set(model.state_dict()): model.load_state_dict(sd["model"], strict=False) print(f"init-from {path} (weights only, mtp heads fresh)", flush=True) else: model.load_state_dict(sd["model"]) print(f"init-from {path} (weights only, fresh schedule)", flush=True) if args.resume: path = latest_ckpt(args.resume) if path: sd = torch.load(path, map_location="cpu") model.load_state_dict(sd["model"]) opt.load_state_dict(sd["opt"]) step, best_val = sd["step"], sd.get("best_val", float("inf")) resume_from = step print(f"resumed {path} at step {step}", flush=True) train_tokens = load_bin(args.data) val_tokens = load_bin(args.val) offsets = make_batches(train_tokens, args.batch, args.seq, rng) # shuffle windows (seed-fixed): every batch is a random domain mix, so a # hard slice of the corpus can never dominate a whole batch offsets = offsets[torch.randperm(len(offsets), generator=torch.Generator().manual_seed(args.seed))] model.train() total_params = model.num_params() print(f"params: {total_params/1e6:.2f}M train_tokens: {len(train_tokens):,} " f"steps_per_epoch: {len(offsets)//args.batch}", flush=True) def save(path, tag=""): payload = { "model": model.state_dict(), "opt": opt.state_dict(), "step": step, "best_val": best_val, "config": cfg.__dict__, "args": vars(args), "tag": tag, } # Never expose a partially written checkpoint to resume logic. tmp = path.with_suffix(path.suffix + ".tmp") torch.save(payload, str(tmp)) os.replace(tmp, path) prompts = ["Once upon a time,", "The little girl wanted to", "In the dark forest,"] target_steps = args.total_steps if args.total_steps is not None else resume_from + args.steps if target_steps < resume_from: raise ValueError(f"total target step {target_steps} is before resume step {resume_from}") t0 = time.time() nan_streak = 0 while step < target_steps: step += 1 x, y = sample_batch(train_tokens, offsets, step - 1, args.batch, args.seq) opt.zero_grad(set_to_none=True) def _fwd(): if args.mtp: logits, aux = model.forward_mtp(x) loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), y.reshape(-1)) for k, a in enumerate(aux): off = k + 2 if args.seq > off: loss = loss + 0.1 * F.cross_entropy( a[:, :-off].reshape(-1, a.size(-1)), x[:, off:].reshape(-1)) return loss logits = model(x) return F.cross_entropy(logits.reshape(-1, logits.size(-1)), y.reshape(-1)) if args.bf16: with torch.autocast("cpu", dtype=torch.bfloat16): loss = _fwd() else: loss = _fwd() if not torch.isfinite(loss): # never let a single bad batch poison the weights: skip the optimizer # step entirely so neither weights nor LR schedule are mutated nan_streak += 1 print(f" !! step {step} non-finite loss {loss.item():.3e}; skipping step", flush=True) if args.nan_rollback > 0 and nan_streak >= args.nan_rollback: rollback = latest_ckpt(ckpt_dir) if rollback is not None: sd = torch.load(rollback, map_location="cpu") model.load_state_dict(sd["model"]) opt.load_state_dict(sd["opt"]) step, best_val = sd["step"], sd.get("best_val", float("inf")) nan_streak = 0 print(f" !! {args.nan_rollback} consecutive non-finite steps; " f"rolled back to {rollback} at step {step}", flush=True) else: nan_streak = 0 print(" !! no checkpoint to roll back to; continuing", flush=True) continue nan_streak = 0 loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) lr = args.lr if step <= args.warmup else args.min_lr + 0.5 * (args.lr - args.min_lr) * ( 1 + math.cos(math.pi * (step - args.warmup) / max(1, target_steps - args.warmup))) for g in opt.param_groups: g["lr"] = lr opt.step() if step % args.log_every == 0 or step == 1: dt = time.time() - t0 toks = args.batch * args.seq print(f"step {step}/{target_steps} loss {loss.item():.4f} lr {lr:.2e} " f"{toks*args.log_every/dt:.0f} tok/s", flush=True) t0 = time.time() if step % args.eval_every == 0: vl = evaluate(model, val_tokens, args.batch, args.seq, args.val_batches, rng, bf16=args.bf16) if vl < best_val: best_val = vl save(ckpt_dir / "model_best.pt", tag="best") save(ckpt_dir / f"{SAVE_PREFIX}_{step}.pt", tag=f"step{step}") print(f" [eval step {step}] val_loss {vl:.4f} (best {best_val:.4f})", flush=True) for pr in prompts: txt = sample(model, tok, pr, args.sample_max_new) print(f" [gen] {pr} {txt}", flush=True) model.train() if step % args.save_every == 0: save(ckpt_dir / f"{SAVE_PREFIX}_{step}.pt", tag=f"step{step}") print(f" saved {ckpt_dir}/{SAVE_PREFIX}_{step}.pt", flush=True) save(ckpt_dir / f"{SAVE_PREFIX}_{step}.pt", tag="final") print("done", flush=True) if __name__ == "__main__": main()