| """Metacognitive halting module for LFM3 (an internal component, the rest of the model is untouched). |
| |
| The output row of </think> is untied from the input embedding and extended by a learned vector D: |
| logit_</think>(h) = (e_</think> + D) . h h = final normalised hidden state |
| so the model's own probability of closing its reasoning becomes a learned function of its internal state ("is the |
| answer I would give now already right?"). Nothing else changes: all other logits, all weights, the input embedding of |
| </think>. Integration is native: a GGUF with a separate output.weight runs in stock llama.cpp, zero extra compute. |
| |
| Training signal (no labels from outside the model except the reference answer of each problem): |
| * sentence ends where the model, forced to answer, is right -> close (y=1) |
| * sentence ends where it is still wrong -> keep thinking (y=0) |
| * the position where the model itself closed thinking -> close (y=1), keep it intact |
| * random mid-sentence positions -> keep thinking (y=0), don't disturb |
| Only D (2048 numbers) is optimised, on the exact sampling probability (full-vocabulary softmax at temperature T). |
| |
| extract: halt_module.py extract --rec results/rec/raw.jsonl --items "data/coc/*.jsonl" --exclude data/coc/mini_math.jsonl |
| --out results/halt/feats.pt |
| train: halt_module.py train --feats results/halt/feats.pt --out results/halt/D.pt [--pos-weight 1 --l2 1e-3] |
| export: halt_module.py export --D results/halt/D.pt --scale 1.0 --gguf-in gguf/lfm3-orig-q5km.gguf --gguf-out gguf/halt.gguf |
| """ |
| import argparse |
| import glob |
| import json |
| import os |
| import random |
| import sys |
|
|
| import torch |
|
|
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) |
| THINK_END = 124902 |
| ap = argparse.ArgumentParser() |
| ap.add_argument("cmd", choices=["extract", "train", "export"]) |
| ap.add_argument("--rec", nargs="+", default=[]) |
| ap.add_argument("--items", nargs="+", default=[]) |
| ap.add_argument("--exclude", nargs="*", default=[]) |
| ap.add_argument("--feats") |
| ap.add_argument("--out") |
| ap.add_argument("--model", default="models/LFM3-1.7B") |
| ap.add_argument("--temp", type=float, default=0.6) |
| ap.add_argument("--n-random", type=int, default=12) |
| ap.add_argument("--pos-weight", type=float, default=1.0) |
| ap.add_argument("--l2", type=float, default=1e-3) |
| ap.add_argument("--epochs", type=int, default=400) |
| ap.add_argument("--D") |
| ap.add_argument("--scale", type=float, default=1.0) |
| ap.add_argument("--gguf-in") |
| ap.add_argument("--gguf-out") |
| args = ap.parse_args() |
|
|
|
|
| def extract(): |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
| import transformers.integrations.sdpa_attention as _sdpa |
| _sdpa.use_gqa_in_sdpa = lambda *a, **k: False |
| from grade import grade_one |
| tok = AutoTokenizer.from_pretrained(args.model) |
| model = AutoModelForCausalLM.from_pretrained(args.model, dtype=torch.bfloat16, attn_implementation="sdpa", |
| trust_remote_code=True).cuda().eval() |
| W = model.get_output_embeddings().weight |
| ends = set() |
| for i in range(len(tok)): |
| t = tok.convert_ids_to_tokens(i) |
| if t is not None and (t.endswith(".") or t.endswith("Ċ")): |
| ends.add(i) |
| items = {} |
| for pat in args.items: |
| for f in glob.glob(pat): |
| for l in open(f, encoding="utf-8"): |
| r = json.loads(l) |
| items[r["id"]] = r |
| excl = set() |
| for f in args.exclude: |
| excl |= {json.loads(l)["id"] for l in open(f, encoding="utf-8")} |
| rng = random.Random(0) |
| H, LSE, BASE, Y, KIND, TRAJ = [], [], [], [], [], [] |
| n_traj = 0 |
| for path in args.rec: |
| for l in open(path, encoding="utf-8"): |
| r = json.loads(l) |
| if r["id"] in excl or r["id"] not in items or not r.get("probe_log"): |
| continue |
| it = items[r["id"]] |
| p_ids = tok.apply_chat_template([{"role": "user", "content": it["prompt"]}], tokenize=True, |
| add_generation_prompt=True, return_dict=False) |
| text = r["completion"][len("<think>"):] |
| resp = tok(text, add_special_tokens=False)["input_ids"][:8192] |
| pos, ys, kinds = [], [], [] |
| |
| |
| for p in r["probe_log"]: |
| c = min(p["cut"], len(resp)) |
| while c > 1 and resp[c - 1] not in ends: |
| c -= 1 |
| if c <= 1: |
| continue |
| ok = grade_one({**r, "completion": "<think>x</think>" + p["text"], "n_tokens": 0}, it)["correct"] |
| pos.append(c - 1) |
| ys.append(1.0 if ok else 0.0) |
| kinds.append("cut") |
| if THINK_END in resp: |
| j = resp.index(THINK_END) |
| pos.append(j - 1) |
| ys.append(1.0) |
| kinds.append("natural") |
| mid = [j for j in range(8, min(len(resp), 8192) - 1) if resp[j] not in ends and resp[j] != THINK_END] |
| for j in rng.sample(mid, min(args.n_random, len(mid))): |
| pos.append(j) |
| ys.append(0.0) |
| kinds.append("random") |
| if not pos: |
| continue |
| ids = p_ids + resp |
| with torch.no_grad(): |
| h = model.model(input_ids=torch.tensor([ids], device="cuda"), use_cache=False).last_hidden_state[0] |
| hp = h[[len(p_ids) + q for q in pos]].float() |
| logits = hp @ W.float().t() / args.temp |
| base = logits[:, THINK_END].clone() |
| logits[:, THINK_END] = -1e30 |
| lse = torch.logsumexp(logits, -1) |
| H.append(hp.cpu()); LSE.append(lse.cpu()); BASE.append(base.cpu()) |
| Y += ys; KIND += kinds; TRAJ += [n_traj] * len(pos) |
| n_traj += 1 |
| if n_traj % 25 == 0: |
| print(n_traj, "trajectories", len(Y), "points", flush=True) |
| os.makedirs(os.path.dirname(args.out), exist_ok=True) |
| torch.save({"h": torch.cat(H), "lse": torch.cat(LSE), "base": torch.cat(BASE), "y": torch.tensor(Y), |
| "kind": KIND, "traj": torch.tensor(TRAJ), "temp": args.temp}, args.out) |
| print("saved", args.out, n_traj, "trajectories", len(Y), "points", |
| {k: (sum(1 for x in KIND if x == k), sum(y for x, y in zip(KIND, Y) if x == k)) for k in set(KIND)}) |
|
|
|
|
| def auc(s, y): |
| order = torch.argsort(s) |
| r = torch.empty_like(order, dtype=torch.float) |
| r[order] = torch.arange(1, len(s) + 1, dtype=torch.float) |
| n1 = y.sum() |
| n0 = len(y) - n1 |
| return ((r[y == 1].sum() - n1 * (n1 + 1) / 2) / (n1 * n0)).item() if n1 > 0 and n0 > 0 else float("nan") |
|
|
|
|
| def train(): |
| d = torch.load(args.feats) |
| h, lse, base, y, traj, T = d["h"], d["lse"], d["base"], d["y"], d["traj"], d["temp"] |
| kind = d["kind"] |
| is_cut = torch.tensor([k == "cut" for k in kind]) |
| n_tr = int(traj.max().item()) + 1 |
| g = torch.Generator().manual_seed(0) |
| val_traj = set(torch.randperm(n_tr, generator=g)[: n_tr // 5].tolist()) |
| val = torch.tensor([t.item() in val_traj for t in traj]) |
| D = torch.zeros(h.shape[1], requires_grad=True) |
| opt = torch.optim.Adam([D], lr=1e-2) |
| w = torch.where(y > 0.5, torch.full_like(y, args.pos_weight), torch.ones_like(y)) |
| tr = ~val |
| for ep in range(args.epochs): |
| z = base + (h @ D) / T - lse |
| loss = torch.nn.functional.binary_cross_entropy_with_logits(z[tr], y[tr], weight=w[tr]) + args.l2 * D.pow(2).sum() |
| opt.zero_grad() |
| loss.backward() |
| opt.step() |
| with torch.no_grad(): |
| z = base + (h @ D) / T - lse |
| p = torch.sigmoid(z) |
| p0 = torch.sigmoid(base - lse) |
| for name, m in [("train", tr), ("val", val)]: |
| c = m & is_cut |
| rnd = m & torch.tensor([k == "random" for k in kind]) |
| nat = m & torch.tensor([k == "natural" for k in kind]) |
| print(f"{name}: cut AUC {auc(z[c], y[c]):.3f} (base {auc((base - lse)[c], y[c]):.3f}) " |
| f"p_close at cuts: right {p[c & (y > .5)].mean():.3f} wrong {p[c & (y < .5)].mean():.3f} " |
| f"(base {p0[c].mean():.4f}) | random pos {p[rnd].mean():.5f} (base {p0[rnd].mean():.5f}) | " |
| f"natural end {p[nat].mean():.3f} (base {p0[nat].mean():.3f})") |
| torch.save({"D": D.detach(), "args": vars(args)}, args.out) |
| print("saved", args.out, "|D|", D.norm().item()) |
|
|
|
|
| def export(): |
| import numpy as np |
| import gguf |
| from safetensors.torch import load_file |
| D = torch.load(args.D)["D"] * args.scale |
| sd = {} |
| for f in glob.glob(os.path.join(args.model, "*.safetensors")): |
| sd.update(load_file(f)) |
| E = sd["model.embed_tokens.weight"].float() |
| E[THINK_END] += D |
| r = gguf.GGUFReader(args.gguf_in) |
| arch = r.fields["general.architecture"].contents() |
| w = gguf.GGUFWriter(args.gguf_out, arch) |
| for f in r.fields.values(): |
| if f.name == "general.architecture" or f.name.startswith("GGUF."): |
| continue |
| vt = f.types[0] |
| st = f.types[-1] if vt == gguf.GGUFValueType.ARRAY else None |
| w.add_key_value(f.name, f.contents(), vt, sub_type=st) |
| w.add_string("general.halt_module", f"untied </think> output row + learned D, scale {args.scale}") |
| q = gguf.quants.quantize(E.numpy().astype(np.float32), gguf.GGMLQuantizationType.Q8_0) |
| tensors = [(t.name, t.data, t.tensor_type) for t in r.tensors] + [("output.weight", q, gguf.GGMLQuantizationType.Q8_0)] |
| for name, data, tt in tensors: |
| w.add_tensor_info(name, data.shape, data.dtype, data.nbytes, tt) |
| w.write_header_to_file(); w.write_kv_data_to_file(); w.write_ti_data_to_file() |
| for name, data, tt in tensors: |
| w.write_tensor_data(data) |
| w.close() |
| print("wrote", args.gguf_out) |
|
|
|
|
| {"extract": extract, "train": train, "export": export}[args.cmd]() |
|
|
