train/train_lm.py

"""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()