train/train_sft2.py

"""Guarded SFT for TinyLiquid: chat+SOP+forensic mix with TinyStories retention.

Differs from train_sft.py:
  * supports raw full-loss retention examples ({"raw": text})
  * evals BOTH masked SFT holdout loss AND TinyStories val PPL (coherence guard)
  * keeps best.pt (min sft_val_loss while val_ppl < 90) and best_ppl.pt (min ppl)

Usage:
  .venv/bin/python train/train_sft2.py --base ckpt/nlp --data data/sft_mix_v2.jsonl \
      --ckpt ckpt/v2 --epochs 3 --lr 2e-5
"""
import argparse, json, math, random, 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

USER_T, ASST_T, EOT_T = "<|user|>", "<|assistant|>", "<|endoftext|>"
PERSONA_T = {"analyst": "<|analyst|>", "skeptic": "<|skeptic|>", "spock": "<|analyst|>", "none": ""}
P_IDS = {"analyst": 1, "skeptic": 2, "spock": 1, "none": 0}

def parse_args():
    ap = argparse.ArgumentParser()
    ap.add_argument("--base", default="ckpt/nlp")
    ap.add_argument("--resume", default="", help="resume from latest ckpt in this dir")
    ap.add_argument("--data", default="data/sft_mix_v2.jsonl")
    ap.add_argument("--tok", default="data/tokenizer.json")
    ap.add_argument("--ckpt", default="ckpt/v2")
    ap.add_argument("--val-bin", default="data/valid.bin")
    ap.add_argument("--epochs", type=int, default=3)
    ap.add_argument("--batch", type=int, default=8)
    ap.add_argument("--seq", type=int, default=256)
    ap.add_argument("--lr", type=float, default=2e-5)
    ap.add_argument("--eval-every", type=int, default=25)
    ap.add_argument("--log-every", type=int, default=25)
    ap.add_argument("--ppl-guard", type=float, default=90.0)
    ap.add_argument("--val-batches", type=int, default=2)
    ap.add_argument("--seed", type=int, default=7)
    ap.add_argument("--threads", type=int, default=8)
    return ap.parse_args()

def tokenize_example(tok, ex, seq, u_id, a_id, eot_id):
    if "raw" in ex:
        ids = tok.encode(ex["raw"]).ids + [eot_id]
        x = torch.tensor(ids[:-1], dtype=torch.long)
        y = torch.tensor(ids[1:], dtype=torch.long)
        mask = torch.ones_like(y, dtype=torch.bool)
        return x[:seq], y[:seq], mask[:seq], 0
    persona = PERSONA_T.get(ex.get("persona", "analyst"), PERSONA_T["analyst"])
    p_id = P_IDS.get(ex.get("persona"), 1)
    user_ids = tok.encode(ex["user"]).ids
    asst_ids = tok.encode(ex["assistant"]).ids
    if persona:
        p_ids_ = tok.encode(persona).ids
        ids = p_ids_ + [u_id] + user_ids + [a_id] + asst_ids + [eot_id]
        asst_start = len(p_ids_) + 1 + len(user_ids) + 1
    else:
        ids = [u_id] + user_ids + [a_id] + asst_ids + [eot_id]
        asst_start = 1 + len(user_ids) + 1
    if len(ids) > seq:
        ids = ids[:seq - 1] + [eot_id]
    x = torch.tensor(ids[:-1], dtype=torch.long)
    y = torch.tensor(ids[1:], dtype=torch.long)
    mask = torch.zeros_like(y, dtype=torch.bool)
    mask[asst_start - 1:] = True
    return x[:seq], y[:seq], mask[:seq], p_id

def collate(items, seq):
    xs, ys, ms, ps = [], [], [], []
    for x, y, m, p in items:
        xs.append(F.pad(x, (0, seq - x.shape[0]), value=0))
        ys.append(F.pad(y, (0, seq - y.shape[0]), value=0))
        ms.append(F.pad(m, (0, seq - m.shape[0]), value=False))
        ps.append(p)
    return torch.stack(xs), torch.stack(ys), torch.stack(ms), torch.tensor(ps, dtype=torch.long)

@torch.no_grad()
def val_ppl(model, val_bin, batch=4, seq=64, n_batches=2, seed=0):
    mm = np.memmap(val_bin, dtype=np.uint16, mode="r")
    total, cnt = 0.0, 0
    rng = np.random.RandomState(seed)
    n = (len(mm) - 1) // seq
    for b in range(n_batches):
        s = int(rng.randint(0, n - batch))
        buf = torch.stack([torch.from_numpy(mm[s * seq + i * seq: s * seq + i * seq + seq].astype(np.int64))
                           for i in range(batch)])
        x, y = buf[:, :-1], buf[:, 1:]
        loss = F.cross_entropy(model(x).reshape(-1, 8192), y.reshape(-1))
        total += loss.item() * y.numel(); cnt += y.numel()
    return float(np.exp(total / cnt))

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)
    tok = load_tokenizer(args.tok)
    u_id, a_id, eot_id = tok.token_to_id(USER_T), tok.token_to_id(ASST_T), tok.token_to_id(EOT_T)
    assert None not in (u_id, a_id, eot_id)

    exs = [json.loads(l) for l in open(args.data, encoding="utf-8") if l.strip()]
    rng.shuffle(exs)
    n_eval = min(128, max(8, len(exs) // 12))
    eval_ex, train_ex = exs[:n_eval], exs[n_eval:]
    print(f"train {len(train_ex)} eval {len(eval_ex)}", flush=True)

    base_path = latest_ckpt(args.resume or args.base)
    base = torch.load(base_path, map_location="cpu")
    base_cfg = base.get("config") or CONFIGS["tiny10m"]
    cfg = TinyLiquidConfig(vocab_size=tok.get_vocab_size(),
                           **{k: v for k, v in base_cfg.items() if k != "vocab_size"})
    model = TinyLiquid(cfg)
    model.load_state_dict(base["model"])
    opt = torch.optim.AdamW(model.parameters(), lr=args.lr, betas=(0.9, 0.95), weight_decay=0.05)
    print(f"loaded base {base_path.name}", flush=True)

    out = Path(args.ckpt); out.mkdir(parents=True, exist_ok=True)
    steps_per_epoch = max(1, len(train_ex) // args.batch)
    total_steps = steps_per_epoch * args.epochs

    def make_items(exs_):
        return [tokenize_example(tok, e, args.seq, u_id, a_id, eot_id) for e in exs_]

    eval_items = make_items(eval_ex)

    def run_eval():
        model.eval()
        total, n = 0.0, 0
        for i in range(0, len(eval_items), args.batch):
            x, y, m, p = collate(eval_items[i:i + args.batch], args.seq)
            with torch.no_grad():
                logits = model(x, persona_ids=p).reshape(-1, 8192)
            loss = F.cross_entropy(logits, y.reshape(-1), reduction="none")
            loss = (loss * m.reshape(-1)).sum() / m.sum()
            total += loss.item() * m.sum().item(); n += m.sum().item()
        sft_vl = total / n
        ppl = val_ppl(model, args.val_bin, n_batches=args.val_batches)
        model.train()
        return sft_vl, ppl

    best_guard, best_ppl = float("inf"), float("inf")
    t0 = time.time(); step = 0
    for ep in range(args.epochs):
        rng.shuffle(train_ex)
        items = make_items(train_ex)
        for i in range(0, len(items) - len(items) % args.batch, args.batch):
            step += 1
            x, y, m, p = collate(items[i:i + args.batch], args.seq)
            opt.zero_grad(set_to_none=True)
            logits = model(x, persona_ids=p).reshape(-1, 8192)
            loss = F.cross_entropy(logits, y.reshape(-1), reduction="none")
            loss = (loss * m.reshape(-1)).sum() / m.sum()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()
            if step % args.log_every == 0:
                print(f"step {step}/{total_steps} loss {loss.item():.4f} "
                      f"{args.batch*args.seq*args.log_every/(time.time()-t0):.0f} tok/s", flush=True)
                t0 = time.time()
            if step % args.eval_every == 0:
                sft_vl, ppl = run_eval()
                try:
                    sp = tok.encode("<|analyst|><|user|>Find discrepancies between: Account A: The meeting ended at 11am. Account B: The meeting ended at noon.<|assistant|>").ids
                    with torch.no_grad():
                        sout = tok.decode(model.generate(tok, sp, persona_id=1, max_new=50, temperature=0.35,
                                                         top_k=20, repetition_penalty=1.25,
                                                         no_repeat_ngram_size=4)[len(sp):]).replace("\n", " ").strip()[:180]
                    print(f"    sample: {sout}", flush=True)
                except Exception:
                    pass
                tag = ""
                if ppl < args.ppl_guard and sft_vl < best_guard:
                    best_guard = sft_vl
                    torch.save({"model": model.state_dict(), "step": step, "config": cfg.__dict__}, out / "best.pt")
                    tag += " [new best]"
                if ppl < best_ppl:
                    best_ppl = ppl
                    torch.save({"model": model.state_dict(), "step": step, "config": cfg.__dict__}, out / "best_ppl.pt")
                    tag += " [best ppl]"
                torch.save({"model": model.state_dict(), "step": step, "config": cfg.__dict__}, out / f"model_{step}.pt")
                print(f"  [eval {step}] sft_val_loss {sft_vl:.4f} val_ppl {ppl:.2f}{tag}", flush=True)
    torch.save({"model": model.state_dict(), "step": step, "config": cfg.__dict__}, out / "model_final.pt")
    print(f"done -> {out}  best_guard_sft_loss={best_guard:.4f} best_ppl={best_ppl:.2f}", flush=True)

if __name__ == "__main__":
    main()