diff --git a/CMakeLists.txt b/CMakeLists.txt index 610142905..d3ca1d191 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -823,6 +823,12 @@ add_executable(laya_eval_test tests/laya_eval_test.cpp) target_link_libraries(laya_eval_test PRIVATE onebit_laya nlohmann_json::nlohmann_json) add_test(NAME laya_eval COMMAND laya_eval_test) add_dependencies(onebit laya_eval_test) +# engine#1313 item 3: the head-only fine-tune over --dump-encoder dumps, on a synthetic checkpoint and dump (no model, +# no device). It needs torch + numpy, which the engine build does not, so a box without them skips (77). +add_python_test(laya_finetune ${CMAKE_SOURCE_DIR}/tests/laya_finetune_test.py) +if(Python3_Interpreter_FOUND) + set_tests_properties(laya_finetune PROPERTIES SKIP_RETURN_CODE 77) +endif() # Crafted Laya checkpoints the reader and the scorer must refuse (written at run time); runs in CI. add_executable(laya_bounds_test tests/laya_bounds_test.cpp) target_link_libraries(laya_bounds_test PRIVATE onebit_laya nlohmann_json::nlohmann_json) diff --git a/docs/laya-models.md b/docs/laya-models.md index 6da158d8b..08af7847b 100644 --- a/docs/laya-models.md +++ b/docs/laya-models.md @@ -113,8 +113,8 @@ numbers a demo HUD reads; every one was measured by serve or reported by the bac `args`. Linux only. - Laya's zero-shot answer to a new question is a starting point, not a measured router: the base checkpoints are near chance on decision types they were not tuned for (Laya's own README). The game - table's accuracy is measured below, and a fine-tune on the game's own logged decisions is the plan for - raising it (engine #1313). + table's accuracy is measured below, and a head-only fine-tune on the game's own logged decisions + (`tools/laya_finetune/`, below) raises its confidence on them (engine #1313). ## Measuring and fine-tuning the router @@ -173,8 +173,51 @@ Measured on this box (`typed-decisions/`, the 20 game steps of the measurement b | a torch copy of the head over the dump | reproduces the engine's probabilities within 2.4e-7 (logits 1.3e-6) on all 40 | | option order | mean \|p(reflex) in one order - the other\| 0.065 (max 0.171); 1 of 20 decisions flips with the order | +### Fine-tuning the head (`tools/laya_finetune/`) + +```sh +# 1. labelled cases -> the head's input, every option order (see the join above for labels) +1bit route --laya-models table.json --laya-model DIR --dump-encoder train.jsonl --out dump-train/ --rotate +1bit route --laya-models table.json --laya-model DIR --dump-encoder dev.jsonl --out dump-dev/ --rotate +# 2. head-only fine-tune on the CPU (torch + numpy): the encoder is frozen, its output is in the dump +tools/laya_finetune/train_head.py --model DIR --train dump-train --dev dump-dev --out NEW_DIR [--epochs 4] [--lr 2e-5] +# 3. before / after on held-out cases, by the path serve takes, and the same numbers from the per-case files +1bit route --laya-models table.json --laya-model DIR --eval test.jsonl --per-case before.jsonl +1bit route --laya-models table.json --laya-model NEW_DIR --eval test.jsonl --per-case after.jsonl +tools/laya_finetune/metrics.py before.jsonl after.jsonl +``` + +`train_head.py` trains `type_emb`, the two head layers and `scorer` (`head.py` is the head of `laya/scorer.cpp` in +torch) on Laya's proper reward (log score + 0.5 spherical score), class-weighted, with early stopping on dev log loss. +It then refits the `choice:k` temperature on dev, and keeps the checkpoint's own temperature when the refit lands on +the grid's edge: a separable dev set says nothing about how sure to be. `NEW_DIR` is a copy of `DIR` with only the head +tensors' bytes replaced (same names, dtypes, shapes and offsets), the temperature in `rl_agent_config.json`, and +`finetune.json` (data, seed, epochs, dev before and after, command line). It refuses to overwrite an existing `NEW_DIR`. +Serve takes it unchanged: `--laya-model NEW_DIR`. + +Measured on the game (engine #1313, the private `laya-plays` harness): labels are the step kind each logged decision +was made for. Episodes were held out whole: `test` = the hand-labelled steps, `test2` and `test3` = later episodes +the fine-tune never saw. + +| set | rows | correct, base / fine-tuned | log loss, base / fine-tuned | +|---|---|---|---| +| test | 39 | 33 / 33 | 0.543 / 1.395 | +| test2 | 110 | 107 / 106 | 0.478 / 0.294 | +| test3 | 56 | 55 / 55 | 0.475 / 0.188 | + +The fine-tune keeps accuracy and makes the router much surer on the game's own steps (log loss 0.48 -> 0.19-0.29). On the +hand-written `test` set it is overconfident on its misses, so read the table row by row, not as one number. A +second label tried in the harness, "plan" for a reflex step whose game action failed, taught the head nothing: 0 of 10 +held-out failed-reflex rows moved to plan, because those failures come from the plan and are not visible in the +reflex step's own state. That label needs a planner's view of the same state, not the reflex prompt. + ## Tests +`tests/laya_finetune_test.py` (ctest `laya_finetune`, no checkpoint; skips without torch): `train_head.py` on a +synthetic checkpoint (the head's real tensor names at d 32) and a dump whose label sits on the right option's marker +row. Dev log loss must fall and accuracy must not; every non-head byte of the new `model.safetensors` must be +identical, the head must change, the temperature must land in its bucket, an existing `--out` must be refused; and +`metrics.py` must agree with a hand-computed `--per-case` file. `tests/laya_models_test.cpp` (ctest `laya_models`, no checkpoint): parsing, the question, the pick and its fallback, aliases, shared backends, the size gate, the refusals, and the shipped game table. `tests/laya_eval_test.cpp` (ctest `laya_eval`, no checkpoint): the eval's arithmetic on hand-computed cases. diff --git a/tests/laya_finetune_test.py b/tests/laya_finetune_test.py new file mode 100755 index 000000000..6323e2885 --- /dev/null +++ b/tests/laya_finetune_test.py @@ -0,0 +1,189 @@ +#!/usr/bin/env python3 +# Copyright 2026 bong-water-water-bong +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""tools/laya_finetune on synthetic data (engine #1313): no checkpoint, no game, no device. + +A tiny checkpoint with the real head's tensor names (d 32, the 16-head layout of laya/scorer.cpp) and one encoder +tensor the fine-tune must not touch, plus `--dump-encoder`-shaped dumps whose label is carried by the marker row of +the right option. Checked: + 1. train_head.py lowers dev log loss and does not lower dev accuracy; + 2. every byte of the new model.safetensors outside the head tensors is identical (header, encoder tensor), the + head tensors changed, and the temperature landed in rl_agent_config.json's bucket; + 3. it refuses to overwrite an existing --out; + 4. metrics.py on a hand-computed --per-case file. +Exit 77 (ctest skip) when torch or numpy is missing: the tool needs them, the engine build does not. +""" +import json +import math +import os +import shutil +import struct +import subprocess +import sys +import tempfile + +try: + import numpy as np + import torch # noqa: F401 +except ImportError: + print("SKIP: torch/numpy not installed") + sys.exit(77) + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +TOOL = os.path.join(ROOT, "tools", "laya_finetune") +sys.path.insert(0, TOOL) +import metrics # noqa: E402 + +D, F_, DH = 32, 64, 32 +fails = 0 + + +def check(ok, what): + global fails + print(("ok " if ok else "FAIL ") + what) + fails += 0 if ok else 1 + + +def write_safetensors(path, tensors): + header, blobs, off = {}, [], 0 + for name, a in tensors.items(): + raw = a.astype(np.float32).tobytes() + header[name] = {"dtype": "F32", "shape": list(a.shape), "data_offsets": [off, off + len(raw)]} + blobs.append(raw) + off += len(raw) + h = json.dumps(header).encode() + h += b" " * (-len(h) % 8) + with open(path, "wb") as f: + f.write(struct.pack(" {a['nll']:.4f}") + check(a["acc"] >= b["acc"], f"dev accuracy does not fall: {b['acc']:.3f} -> {a['acc']:.3f}") + check(a["acc"] >= 0.9, f"the synthetic boundary is learned: dev accuracy {a['acc']:.3f} >= 0.9") + check(meta["train_sequences"] == 48 and meta["class_counts"] == {"reflex": 24, "plan": 24}, + "finetune.json records the data it trained on") + + base0, hdr0 = tensor_ranges(f"{ck}/model.safetensors") + base1, hdr1 = tensor_ranges(f"{out}/model.safetensors") + raw0, raw1 = open(f"{ck}/model.safetensors", "rb").read(), open(f"{out}/model.safetensors", "rb").read() + check(hdr0 == hdr1 and base0 == base1 and len(raw0) == len(raw1), "header, offsets and size are unchanged") + changed, frozen_same = 0, True + for name, info in hdr0.items(): + lo, hi = info["data_offsets"] + same = raw0[base0 + lo:base0 + hi] == raw1[base1 + lo:base1 + hi] + if name.startswith(("head.", "type_emb.", "scorer.")): + changed += 0 if same else 1 + elif not same: + frozen_same = False + check(frozen_same, "every non-head tensor (encoder, temperature) is byte-identical") + check(changed > 0, f"the head tensors were rewritten ({changed} changed)") + cfg = json.load(open(f"{out}/rl_agent_config.json")) + check("choice:2" in cfg.get("temperature_by_options", {}), "the temperature is written to the choice:2 bucket") + + r2 = subprocess.run(cmd, capture_output=True, text=True) + check(r2.returncode != 0 and "refusing to overwrite" in (r2.stdout + r2.stderr), "an existing --out is refused") + + # metrics.py, by hand: 3 cases, 2 right; p(true) 0.9, 0.6, 0.3 + pc = os.path.join(work, "per.jsonl") + with open(pc, "w") as f: + for lab, cho, p, ms in (("reflex", "reflex", 0.9, 100), ("plan", "plan", 0.6, 200), ("plan", "reflex", 0.3, 300)): + other = "plan" if lab == "reflex" else "reflex" + f.write(json.dumps({"label": lab, "chosen": cho, "probabilities": {lab: p, other: 1 - p}, + "decision_ms": ms}) + "\n") + m = metrics.m(pc) + ll = -(math.log(0.9) + math.log(0.6) + math.log(0.3)) / 3 + # top p 0.9 (right), 0.6 (right), 0.7 (wrong) -> bins 9, 6, 7, each |p - correct|: 0.1, 0.4, 0.7 + ece = (0.1 + 0.4 + 0.7) / 3 + check(m["correct"] == 2 and m["acc"] == round(2 / 3, 4), "metrics: accuracy 2/3") + check(m["log_loss"] == round(ll, 4), f"metrics: log loss {m['log_loss']} == {round(ll, 4)}") + check(m["ece"] == round(ece, 4), f"metrics: ECE {m['ece']} == {round(ece, 4)}") + check(m["ms_p50"] == 200 and m["min_p_true"] == 0.3, "metrics: decision p50 and min p(true)") + finally: + shutil.rmtree(work, ignore_errors=True) + + +if __name__ == "__main__": + os.makedirs(os.path.expanduser("~/.cache/scratch-tests"), exist_ok=True) + main() + print("all laya_finetune checks passed" if fails == 0 else f"{fails} laya_finetune check(s) FAILED") + sys.exit(1 if fails else 0) diff --git a/tools/laya_finetune/head.py b/tools/laya_finetune/head.py new file mode 100644 index 000000000..d0de081ac --- /dev/null +++ b/tools/laya_finetune/head.py @@ -0,0 +1,118 @@ +#!/usr/bin/env python3 +# Copyright 2026 bong-water-water-bong +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Laya's decision head in plain torch, over the encoder states `1bit route --dump-encoder` writes. + +It is the head of laya/scorer.cpp (and of DecisionModel upstream): h + type_emb[qtype], two pre-norm +TransformerEncoderLayers (d 1024, 16 heads, ff 4096, relu, key padding), the option markers, and the scorer +LayerNorm -> Linear -> GELU(erf) -> Linear(1); probabilities are softmax(logit / T) with T from +temperature_by_options. The encoder is frozen: only these tensors are trained. +""" +import json +import struct + +import numpy as np +import torch +import torch.nn.functional as F + +HEAD_PREFIXES = ("head.", "type_emb.", "scorer.") +DT = {"F32": np.float32, "F16": np.float16, "BF16": None} + + +def read_safetensors(path, want=lambda name: True): + with open(path, "rb") as f: + n = struct.unpack("= 0] + nll, ok = 0.0, 0 + with torch.no_grad(): + for lg, y, cs in logits_of(w, dump, idx): + for b, c in enumerate(cs): + k = len(c["markers"]) + t = temp if temp is not None else H.temperature(w, cfg, c["qtype"], k) + lp = torch.log_softmax(lg[b, :k] / t, -1) + nll -= lp[y[b]].item() + ok += int(lp.argmax().item() == y[b].item()) + return nll / max(1, len(idx)), ok / max(1, len(idx)), len(idx) + + +def refit_temperature(w, dump): + """The temperature with the lowest NLL on dump, over a 0.2..4.0 grid; the head runs once.""" + idx = [i for i, c in enumerate(dump.cases) if c["label"] >= 0] + rows = [] + with torch.no_grad(): + for lg, y, cs in logits_of(w, dump, idx): + for b, c in enumerate(cs): + rows.append((lg[b, :len(c["markers"])].clone(), int(y[b]))) + grid = [round(0.2 + 0.02 * i, 2) for i in range(191)] # 0.2 .. 4.0 + + def nll(t): + return -sum(torch.log_softmax(lg / t, -1)[y].item() for lg, y in rows) / len(rows) + + return min(grid, key=nll) + + +def proper_reward(logp, y): + """Laya's proper scoring reward: log score + 0.5 spherical score (both strictly proper).""" + p = logp.exp() + log_score = logp.gather(-1, y[:, None])[:, 0] + sph = p.gather(-1, y[:, None])[:, 0] / p.norm(dim=-1).clamp_min(1e-9) + return log_score + 0.5 * sph + + +def write_checkpoint(src_dir, out_dir, w, new_temp, k_bucket): + if os.path.exists(out_dir): + raise SystemExit(f"{out_dir} exists; refusing to overwrite") + shutil.copytree(src_dir, out_dir) + path = f"{out_dir}/model.safetensors" + with open(path, "r+b") as f: + n = struct.unpack("= 0] + if not lab: + raise SystemExit("no labelled sequences in --train") + ks = {len(train.cases[i]["markers"]) for i in lab} + qts = {train.cases[i]["qtype"] for i in lab} + if len(ks) != 1 or qts != {0}: + raise SystemExit("one choice question with one option count per dump is supported") + k = ks.pop() + kb = H.bucket(0, k) + # class weights over the labels by class name (rotations move the index, not the class) + names = [train.cases[i]["options"][train.cases[i]["label"]] for i in lab] + freq = {n: names.count(n) for n in set(names)} + cw = {n: len(names) / (len(freq) * c) for n, c in freq.items()} + + params = [v for n, v in w.items() if n != "temperature"] + opt = torch.optim.AdamW(params, lr=a.lr, weight_decay=a.wd) + T0 = H.temperature(w, cfg, 0, k) + before = evaluate(w, cfg, dev) if dev else None + best = (before[0] if before else math.inf, {n: v.detach().clone() for n, v in w.items()}, 0) + log = [] + for ep in range(1, a.epochs + 1): + random.shuffle(lab) + tot, nb = 0.0, 0 + for s in range(0, len(lab), a.bs): + idx = lab[s:s + a.bs] + h, pad, mk, qt, y, cs = train.batch(idx) + lg = H.forward(w, h, pad, mk, qt)[:, :k] + logp = torch.log_softmax(lg / T0, -1) + wt = torch.tensor([cw[c["options"][c["label"]]] for c in cs]) + loss = -(proper_reward(logp, y) * wt).sum() / wt.sum() + opt.zero_grad() + loss.backward() + torch.nn.utils.clip_grad_norm_(params, 1.0) + opt.step() + tot += loss.item() + nb += 1 + d = evaluate(w, cfg, dev) if dev else None + log.append({"epoch": ep, "train_loss": tot / nb, "dev_nll": d[0] if d else None, "dev_acc": d[1] if d else None}) + print(json.dumps(log[-1]), flush=True) + if d is None or d[0] < best[0]: + best = (d[0] if d else -ep, {n: v.detach().clone() for n, v in w.items()}, ep) + if dev: + for n in w: + w[n].data.copy_(best[1][n]) + # temperature refit on dev (or train when there is no dev) + ref = dev or train + temp = refit_temperature(w, ref) + refit_kept = True + if temp in (0.2, 4.0): + # A refit at the grid's edge means the dev set is (nearly) separable and says nothing about how sure to be. + # ft1 (refit 0.2 on a template-separable dev) made its one out-of-template miss confidently wrong, p(true) + # 0.446 -> 0.001. Keep the checkpoint's own temperature instead. + print(f"warning: the refit temperature {temp} is at the grid's edge; keeping {T0:.4f}", flush=True) + temp, refit_kept = float(T0), False + after = evaluate(w, cfg, dev, temp) if dev else None + write_checkpoint(a.model, a.out, w, temp, kb) + meta = {"from": os.path.abspath(a.model), "train": os.path.abspath(a.train), + "dev": os.path.abspath(a.dev) if a.dev else None, "train_sequences": len(lab), "class_counts": freq, + "class_weights": cw, "seed": a.seed, "epochs": a.epochs, "best_epoch": best[2], "lr": a.lr, + "weight_decay": a.wd, "batch": a.bs, "temperature_bucket": kb, "temperature_before": T0, + "temperature_after": temp, "temperature_refit_kept": refit_kept, "dev_before": {"nll": before[0], "acc": before[1], "n": before[2]} if before else None, + "dev_after": {"nll": after[0], "acc": after[1], "n": after[2]} if after else None, "log": log, + "seconds": round(time.time() - t_start, 1), "argv": sys.argv} + json.dump(meta, open(f"{a.out}/finetune.json", "w"), indent=1) + print(json.dumps({k2: meta[k2] for k2 in ("best_epoch", "temperature_before", "temperature_after", + "dev_before", "dev_after", "seconds")})) + + +if __name__ == "__main__": + main()