Buckets:
| """LDR: latent diffusion reasoner. The backbone runs ONCE; thinking = K refinement steps on a small set of latent vectors. | |
| context = frozen P1 encoder (S1 trained on procedural single questions, 88.7%) -> last-layer states H, option states | |
| Z_0 = learned init + noise (N latents, d = 1024) | |
| Z_k = R(Z_{k-1}, k, H) R = weight-tied block: self-attn over Z, cross-attn to H, MLP (x n_layers) | |
| decision = P1's readout on (option state + head(option state -> Z_k)), head output zero-init => step 0 == P1 | |
| Training: deep supervision (decision loss after every step k = 1..K), optional per-step noise (diffusion-style). | |
| Cost per step: N latents through n_layers small blocks (cross-attending to cached H) -- not a backbone pass. | |
| """ | |
| import argparse | |
| import json | |
| import math | |
| import os | |
| import random | |
| import sys | |
| import time | |
| import numpy as np | |
| import torch | |
| import torch.nn.functional as F | |
| from torch import nn | |
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) | |
| from common import PROC, Renderer, load_canonical, tokenizer_and_mapping # noqa: E402 | |
| from model import collate, item_logits # noqa: E402 | |
| from train import micro_batches # noqa: E402 | |
| from shared_yesno_supervised import supervised_terms # noqa: E402 | |
| class Attn(nn.Module): | |
| def __init__(self, d, h): | |
| super().__init__() | |
| self.h = h | |
| self.q, self.k, self.v, self.o = (nn.Linear(d, d, bias=False) for _ in range(4)) | |
| def forward(self, x, ctx, mask=None): | |
| B, N, d = x.shape | |
| M = ctx.shape[1] | |
| q = self.q(x).view(B, N, self.h, -1).transpose(1, 2) | |
| k = self.k(ctx).view(B, M, self.h, -1).transpose(1, 2) | |
| v = self.v(ctx).view(B, M, self.h, -1).transpose(1, 2) | |
| a = F.scaled_dot_product_attention(q, k, v, attn_mask=None if mask is None else mask[:, None, None, :]) | |
| return self.o(a.transpose(1, 2).reshape(B, N, d)) | |
| class Block(nn.Module): | |
| def __init__(self, d, h): | |
| super().__init__() | |
| self.n1, self.n2, self.n3 = nn.LayerNorm(d), nn.LayerNorm(d), nn.LayerNorm(d) | |
| self.sa, self.ca = Attn(d, h), Attn(d, h) | |
| self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d)) | |
| def forward(self, z, ctx, mask): | |
| y = self.n1(z) | |
| z = z + self.sa(y, y) | |
| z = z + self.ca(self.n2(z), ctx, mask) | |
| return z + self.mlp(self.n3(z)) | |
| class Reasoner(nn.Module): | |
| def __init__(self, d=1024, n_lat=32, n_layers=2, heads=16, max_steps=64): | |
| super().__init__() | |
| self.z0 = nn.Parameter(torch.randn(n_lat, d) * 0.02) | |
| self.step_emb = nn.Embedding(max_steps + 1, d) | |
| nn.init.zeros_(self.step_emb.weight) | |
| self.ctx_norm = nn.LayerNorm(d) | |
| self.blocks = nn.ModuleList([Block(d, heads) for _ in range(n_layers)]) | |
| # decision head: option states attend to latents; zero-init output => P1 at init | |
| self.hq, self.hk = nn.LayerNorm(d), nn.LayerNorm(d) | |
| self.head = Attn(d, heads) | |
| nn.init.zeros_(self.head.o.weight) | |
| def init(self, B, noise): | |
| z = self.z0[None].expand(B, -1, -1) | |
| return z + noise * torch.randn_like(z) if noise else z | |
| def step(self, z, k, ctx, mask, noise=0.0): | |
| if noise: | |
| z = z + noise * torch.randn_like(z) | |
| z = z + self.step_emb.weight[min(k, self.step_emb.num_embeddings - 1)] | |
| for b in self.blocks: | |
| z = b(z, ctx, mask) | |
| return z | |
| def delta(self, h_opt, z): | |
| return self.head(self.hq(h_opt), self.hk(z)) | |
| def encode(p1, batch): | |
| """Frozen P1: last-layer pre-norm states H and head-adjusted option states (P1's own pre-readout state).""" | |
| ids, attn = batch["ids"], batch["attn"] | |
| x = p1.bb.embed_tokens(ids) | |
| pos = torch.arange(ids.shape[1], device=ids.device)[None] | |
| h = p1.layers(x, attn[:, None, None, :], p1.bb.rotary_emb(x, pos)) | |
| b = torch.arange(h.shape[0], device=h.device)[:, None] | |
| h_opt = h[b, batch["opt_idx"]] | |
| h_opt = h_opt + p1.head(h_opt, batch["opt_mask"], h_opt, batch["opt_mask"]).to(h_opt.dtype) | |
| return h, h_opt | |
| def run_steps(p1, R, batch, K, noise0, step_noise, reads): | |
| h, h_opt = encode(p1, batch) | |
| ctx = R.ctx_norm(h.float()) | |
| z = R.init(h.shape[0], noise0) | |
| outs = {} | |
| if 0 in reads: | |
| outs[0] = p1.readout(h_opt) | |
| for k in range(1, K + 1): | |
| z = R.step(z, k, ctx, batch["attn"], step_noise) | |
| if k in reads: | |
| outs[k] = p1.readout(h_opt.float() + R.delta(h_opt.float(), z)) | |
| return outs | |
| def load_p1(path, device, mp): | |
| from eval import build_model | |
| p1 = build_model(path, 0, True, device, mp) | |
| for p in p1.parameters(): | |
| p.requires_grad_(False) | |
| return p1.eval() | |
| def train(a): | |
| device = torch.device("cuda") | |
| torch.manual_seed(a.seed) | |
| random.seed(a.seed) | |
| tok, mp, mid, mt = tokenizer_and_mapping() | |
| rend = Renderer(tok, mid, mt, 0) | |
| items, _ = load_canonical(a.train_file, rend, 10**9, a.train_limit) | |
| if a.skip: # never train on the encoder's own training items if requested | |
| items = items[a.skip:] | |
| p1 = load_p1(a.p1, device, mp) | |
| p1.compile_layers() | |
| R = Reasoner(n_lat=a.n_lat, n_layers=a.n_layers).to(device) | |
| print(json.dumps(dict(event="start", items=len(items), reasoner_params=sum(p.numel() for p in R.parameters()), | |
| K=a.k_train, n_lat=a.n_lat, n_layers=a.n_layers, epochs=a.epochs)), flush=True) | |
| opt = torch.optim.AdamW(R.parameters(), lr=a.lr, weight_decay=0.01, fused=True) | |
| steps_per_epoch = len(items) // a.batch | |
| total = steps_per_epoch * a.epochs | |
| warm = max(1, int(0.05 * total)) | |
| log = open(os.path.join(a.out, "train_log.jsonl"), "a") | |
| t0, step = time.time(), 0 | |
| rng = random.Random(a.seed) | |
| for ep in range(a.epochs): | |
| order = list(range(len(items))) | |
| rng.shuffle(order) | |
| for s in range(steps_per_epoch): | |
| f = (step + 1) / warm if step < warm else 0.5 * (1 + math.cos(math.pi * (step - warm) / max(1, total - warm))) | |
| for g in opt.param_groups: | |
| g["lr"] = a.lr * max(f, 0.02) | |
| batch_items = [items[i] for i in order[s * a.batch:(s + 1) * a.batch]] | |
| tot = {k: 0.0 for k in range(1, a.k_train + 1)} | |
| for mb in micro_batches(batch_items, 1, a.budget, 128): | |
| b = collate(mb, device, tok.pad_token_id) | |
| with torch.autocast("cuda", dtype=torch.bfloat16): | |
| outs = run_steps(p1, R, b, a.k_train, a.noise0, a.step_noise, set(range(1, a.k_train + 1))) | |
| loss = 0.0 | |
| for k, sc in outs.items(): | |
| lk = sum(supervised_terms(lg.float(), torch.tensor([x["target"]], device=device), x["primitive"], 0.25)["loss"] | |
| for lg, x in zip(item_logits(sc, b), mb)) | |
| tot[k] += float(lk) | |
| loss = loss + lk | |
| loss = loss / (len(batch_items) * len(outs)) | |
| loss.backward() | |
| gn = torch.nn.utils.clip_grad_norm_(R.parameters(), 1.0) | |
| opt.step() | |
| opt.zero_grad(set_to_none=True) | |
| step += 1 | |
| if step % 20 == 0 or step == total: | |
| rec = dict(step=step, total=total, epoch=ep + 1, grad_norm=round(float(gn), 3), elapsed=round(time.time() - t0, 1), | |
| **{"loss_k%d" % k: round(v / len(batch_items), 4) for k, v in tot.items() if k in (1, 2, 4, a.k_train)}, | |
| head_o=round(R.head.o.weight.norm().item(), 4), peak_gib=round(torch.cuda.max_memory_allocated() / 2**30, 1)) | |
| log.write(json.dumps(rec) + "\n") | |
| log.flush() | |
| print(json.dumps(rec), flush=True) | |
| torch.save(dict(state=R.state_dict(), cfg=vars(a)), os.path.join(a.out, "reasoner.pt")) | |
| print(json.dumps(dict(event="saved", elapsed=round(time.time() - t0, 1))), flush=True) | |
| def evaluate(a): | |
| device = torch.device("cuda") | |
| tok, mp, mid, mt = tokenizer_and_mapping() | |
| rend = Renderer(tok, mid, mt, 0) | |
| path = PROC + ("/test.jsonl" if a.split == "test" else "/dev.jsonl") | |
| items, _ = load_canonical(path, rend, 10**9) | |
| p1 = load_p1(a.p1, device, mp) | |
| ck = torch.load(os.path.join(a.out, "reasoner.pt"), weights_only=False) | |
| cfg = ck["cfg"] | |
| R = Reasoner(n_lat=cfg["n_lat"], n_layers=cfg["n_layers"]).to(device) | |
| R.load_state_dict(ck["state"]) | |
| R.eval() | |
| reads = sorted(set([0, 1, 2, 3, 4, 6, 8, 12, 16, 24, 32])) | |
| K = max(reads) | |
| preds = {k: [None] * len(items) for k in reads} | |
| order = sorted(range(len(items)), key=lambda i: len(items[i]["ids"])) | |
| chunks, cur, L = [], [], 0 | |
| for i in order: | |
| n = len(items[i]["ids"]) | |
| if cur and (n * (len(cur) + 1) > 65536 or len(cur) >= 64): | |
| chunks.append(cur) | |
| cur = [] | |
| cur.append(i) | |
| if cur: | |
| chunks.append(cur) | |
| t0 = time.time() | |
| for ch in chunks: | |
| mb = [items[i] for i in ch] | |
| b = collate(mb, device, tok.pad_token_id) | |
| with torch.autocast("cuda", dtype=torch.bfloat16): | |
| outs = run_steps(p1, R, b, K, cfg["noise0"] if a.noise_eval else 0.0, 0.0, set(reads)) | |
| for k, sc in outs.items(): | |
| for j, (lg, x) in enumerate(zip(item_logits(sc.float(), b), mb)): | |
| preds[k][ch[j]] = torch.softmax(lg[0].float(), -1).cpu().numpy() | |
| rep = dict(split=a.split, n=len(items), seconds=round(time.time() - t0, 1), by_K={}) | |
| hard = [i for i, x in enumerate(items) if x["kind"] == "hard"] | |
| for k in reads: | |
| corr = {i: int(np.argmax(preds[k][i]) == int(np.argmax(items[i]["target"]))) for i in hard} | |
| by_level = {l: float(np.mean([corr[i] for i in hard if items[i]["level"] == l])) for l in range(5)} | |
| rep["by_K"][str(k)] = dict(all=float(np.mean(list(corr.values()))), by_level=by_level) | |
| print("K=%2d all %.2f | %s" % (k, 100 * rep["by_K"][str(k)]["all"], " ".join("L%d %.2f" % (l, 100 * v) for l, v in by_level.items())), | |
| flush=True) | |
| with open(os.path.join(a.out, "pred_%s.jsonl" % a.split), "w") as f: | |
| for i, x in enumerate(items): | |
| f.write(json.dumps(dict(id="%s|%s" % (x["case_id"], x["view_id"]), suite=x["source"], level=x["level"], | |
| gold_idx=int(np.argmax(x["target"])) if x["kind"] == "hard" else None, | |
| probs={str(k): [round(float(v), 6) for v in preds[k][i]] for k in reads})) + "\n") | |
| json.dump(rep, open(os.path.join(a.out, "report_%s.json" % a.split), "w"), indent=1) | |
| if __name__ == "__main__": | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("mode", choices=["train", "eval"]) | |
| ap.add_argument("--out", required=True) | |
| ap.add_argument("--p1", default="/workspace/think-decide/runs/p1/model.safetensors") | |
| ap.add_argument("--train-file", default=PROC + "/train.jsonl") | |
| ap.add_argument("--train-limit", type=int, default=0) | |
| ap.add_argument("--skip", type=int, default=0) | |
| ap.add_argument("--epochs", type=int, default=1) | |
| ap.add_argument("--batch", type=int, default=64) | |
| ap.add_argument("--lr", type=float, default=3e-4) | |
| ap.add_argument("--budget", type=int, default=65536) | |
| ap.add_argument("--n-lat", type=int, default=32) | |
| ap.add_argument("--n-layers", type=int, default=2) | |
| ap.add_argument("--k-train", type=int, default=8) | |
| ap.add_argument("--noise0", type=float, default=1.0) | |
| ap.add_argument("--step-noise", type=float, default=0.0) | |
| ap.add_argument("--seed", type=int, default=42) | |
| ap.add_argument("--split", choices=["test", "dev"], default="test") | |
| ap.add_argument("--noise-eval", action="store_true") | |
| a = ap.parse_args() | |
| os.makedirs(a.out, exist_ok=True) | |
| train(a) if a.mode == "train" else evaluate(a) | |
Xet Storage Details
- Size:
- 11.8 kB
- Xet hash:
- 6aa685406c215bbc302e4eca840b66b56d8e432f4737556abf0bf271a9209579
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.