File size: 6,865 Bytes
97c39f2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 | """Supervised fine-tune of TinyLiquid for forensic analysis (SOP + scratchpad).
Format per example:
<|persona|><|user|>USER<|assistant|>ASSISTANT<|endoftext|>
Loss is masked to the ASSISTANT segment (including scratchpad markers).
Usage:
.venv/bin/python train/train_sft.py --base ckpt/nlp --data data/sft_forensic.jsonl \
--ckpt ckpt/forensic --epochs 3
"""
import argparse
import json
import math
import random
import time
from pathlib import Path
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 = "<|user|>"
ASST_T = "<|assistant|>"
EOT_T = "<|endoftext|>"
PERSONA_T = {"analyst": "<|analyst|>", "skeptic": "<|skeptic|>", "spock": "<|analyst|>"}
def parse_args():
ap = argparse.ArgumentParser()
ap.add_argument("--base", default="ckpt/nlp", help="dir with pretrain checkpoints")
ap.add_argument("--data", default="data/sft_forensic.jsonl")
ap.add_argument("--tok", default="data/tokenizer.json")
ap.add_argument("--ckpt", default="ckpt/forensic")
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=5e-5)
ap.add_argument("--eval-every", type=int, default=200)
ap.add_argument("--log-every", type=int, default=25)
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, p_ids):
persona = PERSONA_T.get(ex["persona"], PERSONA_T["analyst"])
p_id = p_ids[ex["persona"]] if ex["persona"] in p_ids else p_ids["analyst"]
parts = [tok.encode(persona).ids, [u_id], tok.encode(ex["user"]).ids,
[a_id], tok.encode(ex["assistant"]).ids, [eot_id]]
ids = [i for part in parts for i in part]
if len(ids) > seq: # truncate assistant side
keep = seq - 1
ids = ids[:keep] + [eot_id]
asst_start = len(parts[0]) + 1 + len(parts[2]) + 1 # index of first assistant token
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 # y positions that predict assistant tokens
mask = mask[: x.shape[0]]
return x, y[: x.shape[0]], mask, 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))
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 = tok.token_to_id(USER_T)
a_id = tok.token_to_id(ASST_T)
eot_id = tok.token_to_id(EOT_T)
# persona embedding indices (NOT tokenizer ids): 0=none, 1=analyst, 2=skeptic
p_ids = {"analyst": 1, "skeptic": 2}
assert None not in (u_id, a_id, eot_id), "special tokens missing from tokenizer"
examples = [json.loads(l) for l in open(args.data, encoding="utf-8") if l.strip()]
rng.shuffle(examples)
n_eval = min(128, len(examples) // 10)
eval_ex, train_ex = examples[:n_eval], examples[n_eval:]
print(f"train {len(train_ex)} eval {len(eval_ex)}", flush=True)
base_ckpt = latest_ckpt(args.base)
assert base_ckpt, f"no pretrain checkpoint in {args.base}"
base = torch.load(base_ckpt, map_location="cpu")
config = base.get("config") or CONFIGS["tiny10m"]
cfg = TinyLiquidConfig(vocab_size=tok.get_vocab_size(),
**{k: v for k, v in config.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_ckpt.name} (step {base.get('step', '?')})", flush=True)
out_dir = Path(args.ckpt)
out_dir.mkdir(parents=True, exist_ok=True)
steps_per_epoch = max(1, len(train_ex) // args.batch)
total_steps = steps_per_epoch * args.epochs
model.train()
def make_items(exs):
return [tokenize_example(tok, e, args.seq, u_id, a_id, eot_id, p_ids) for e in exs]
def run_eval():
model.eval()
items = make_items(eval_ex)
total, n = 0.0, 0
for i in range(0, len(items), args.batch):
x, y, m, p = collate(items[i:i + args.batch], args.seq)
with torch.no_grad():
logits = model(x, persona_ids=p)
logits = logits.reshape(-1, logits.size(-1))
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()
model.train()
return total / n
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)
logits = logits.reshape(-1, logits.size(-1))
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:
dt = time.time() - t0
print(f"step {step}/{total_steps} loss {loss.item():.4f} "
f"{args.batch*args.seq*args.log_every/dt:.0f} tok/s", flush=True)
t0 = time.time()
if step % args.eval_every == 0:
vl = run_eval()
print(f" [eval {step}] sft_val_loss {vl:.4f}", flush=True)
torch.save({"model": model.state_dict(), "opt": opt.state_dict(),
"step": step, "config": cfg.__dict__},
str(out_dir / f"model_{step}.pt"))
torch.save({"model": model.state_dict(), "opt": opt.state_dict(),
"step": step, "config": cfg.__dict__}, str(out_dir / "model_final.pt"))
print(f"done -> {out_dir}", flush=True)
if __name__ == "__main__":
main()
|