Buckets:
| """LDR v2: latent diffusion reasoner with answer-belief feedback (TRM-style), random depth, truncated BPTT. | |
| Backbone (frozen) runs ONCE: context states H + option states for every question on the canvas. | |
| Z_0 = z0 + 0.5 N(0, I) (same at train and eval; eval averages 4 noise seeds) | |
| Y_k = LN(h_opt) + W_b [p_k, log p_k] one slot per answer mask, p_k = current beliefs (detached) | |
| Z_k = R(Z_{k-1}; [Z; Y_k], H) weight-tied blocks: self-attn over [Z;Y], cross-attn to H, MLP (no step embedding) | |
| decision_k = readout(h_opt + head(h_opt -> Z_k)) head output zero-init => K=0 == frozen encoder | |
| Train: K ~ U{1..12} per batch; steps 1..K-3 without grad, last 3 with grad; loss = CE(step K) + 0.3 mean CE(other grad steps). | |
| Mode "single": one question per canvas, encoder P1 (+ its head). Mode "multi": all questions of a state (MQ canvas), | |
| encoder MQ, every question supervised; one belief slot per answer mask of every question. | |
| """ | |
| import argparse | |
| import copy | |
| 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 ldr import Attn # noqa: E402 | |
| from shared_yesno_supervised import supervised_terms # noqa: E402 | |
| class Block(nn.Module): | |
| def __init__(self, d, h): | |
| super().__init__() | |
| self.n1, self.n1k, self.n2, self.n3 = nn.LayerNorm(d), 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, y, ymask, ctx, cmask): | |
| kv = torch.cat([z, y], 1) | |
| kvm = torch.cat([torch.ones(z.shape[:2], dtype=torch.bool, device=z.device), ymask], 1) | |
| z = z + self.sa(self.n1(z), self.n1k(kv), kvm) | |
| z = z + self.ca(self.n2(z), ctx, cmask) | |
| return z + self.mlp(self.n3(z)) | |
| class Reasoner(nn.Module): | |
| def __init__(self, d=1024, n_lat=32, n_layers=2, heads=16): | |
| super().__init__() | |
| self.z0 = nn.Parameter(torch.randn(n_lat, d) * 0.02) | |
| self.ctx_norm, self.y_norm = nn.LayerNorm(d), nn.LayerNorm(d) | |
| self.belief = nn.Linear(2, d) | |
| self.blocks = nn.ModuleList([Block(d, heads) for _ in range(n_layers)]) | |
| self.hq, self.hk = nn.LayerNorm(d), nn.LayerNorm(d) | |
| self.head = Attn(d, heads) | |
| nn.init.zeros_(self.head.o.weight) | |
| def y_slots(self, h_opt, p): | |
| return self.y_norm(h_opt) + self.belief(torch.stack([p, torch.log(p.clamp_min(1e-6))], -1)) | |
| def step(self, z, y, ymask, ctx, cmask): | |
| for b in self.blocks: | |
| z = b(z, y, ymask, ctx, cmask) | |
| return z | |
| def delta(self, h_opt, z): | |
| return self.head(self.hq(h_opt), self.hk(z)) | |
| # ------------------------------------------------------------------------------------------------ data | |
| def single_items(path, limit, skip, tok, mid, mt): | |
| items, _ = load_canonical(path, Renderer(tok, mid, mt, 0), 10**9, limit) | |
| out = [] | |
| for x in items[skip:]: | |
| out.append(dict(ids=x["ids"], id="%s|%s" % (x["case_id"], x["view_id"]), level=x["level"], task=x["source"], | |
| qs=[dict(pos=x["pos"], prim=x["primitive"], target=x["target"], kind=x["kind"], is_target=True)])) | |
| return out | |
| def multi_items(split, canon_path, skip, limit, tok, mid, mt, shuffle): | |
| from mq import load_states, prepare | |
| rows = [json.loads(l) for l in open(canon_path)] | |
| rows = rows[skip:skip + limit] if limit else rows[skip:] | |
| target_q = {r["group_id"]: next(iter(r["questions"])) for r in rows} | |
| states = load_states(split, list(target_q), shuffle_questions=shuffle, seed=7) | |
| states, _ = prepare(states, tok, mid, mt, 10**9) | |
| out = [] | |
| for s in states: | |
| out.append(dict(ids=s["ids"], id=s["id"], level=s["level"], task="proc/" + s["task"], | |
| qs=[dict(pos=x["pos"], prim=x["prim"], target=x["target"], qid=x["qid"], | |
| kind="hard" if max(x["target"]) == 1.0 else "soft", | |
| is_target=x["qid"] == target_q.get(s["id"])) for x in s["qs"]])) | |
| return out | |
| def collate(items, device, pad_id): | |
| B, L = len(items), max(len(x["ids"]) for x in items) | |
| slots = [[p for q in x["qs"] for p in q["pos"]] for x in items] | |
| N = max(map(len, slots)) | |
| ids = torch.full((B, L), pad_id, dtype=torch.long) | |
| attn = torch.zeros((B, L), dtype=torch.bool) | |
| oi = torch.zeros((B, N), dtype=torch.long) | |
| om = torch.zeros((B, N), dtype=torch.bool) | |
| for i, x in enumerate(items): | |
| ids[i, :len(x["ids"])] = torch.tensor(x["ids"]) | |
| attn[i, :len(x["ids"])] = True | |
| oi[i, :len(slots[i])] = torch.tensor(slots[i]) | |
| om[i, :len(slots[i])] = True | |
| return dict(ids=ids.to(device), attn=attn.to(device), opt_idx=oi.to(device), opt_mask=om.to(device), items=items) | |
| def per_question(s, items): | |
| """s [B,N] slot scores -> list over items of list over questions of logits [1,n] (noul -> [0, s]).""" | |
| out = [] | |
| for i, x in enumerate(items): | |
| k, qs = 0, [] | |
| for q in x["qs"]: | |
| n = len(q["pos"]) | |
| sc = s[i, k:k + n] | |
| qs.append(torch.stack([sc.new_zeros(()), sc[0]])[None] if q["prim"] == "noul" else sc[None]) | |
| k += n | |
| out.append(qs) | |
| return out | |
| def slot_beliefs(s, items): | |
| """Current belief per slot: choice/score -> option prob within its question; noul -> p(Yes).""" | |
| p = torch.zeros_like(s, dtype=torch.float32) | |
| for i, x in enumerate(items): | |
| k = 0 | |
| for q in x["qs"]: | |
| n = len(q["pos"]) | |
| sc = s[i, k:k + n].float() | |
| p[i, k:k + n] = torch.sigmoid(sc) if q["prim"] == "noul" else torch.softmax(sc, -1) | |
| k += n | |
| return p | |
| # ------------------------------------------------------------------------------------------------ model glue | |
| def encode(enc, batch, use_enc_head): | |
| ids, attn = batch["ids"], batch["attn"] | |
| x = enc.bb.embed_tokens(ids) | |
| pos = torch.arange(ids.shape[1], device=ids.device)[None] | |
| h = enc.layers(x, attn[:, None, None, :], enc.bb.rotary_emb(x, pos)) | |
| b = torch.arange(h.shape[0], device=h.device)[:, None] | |
| h_opt = h[b, batch["opt_idx"]] | |
| if use_enc_head: | |
| h_opt = h_opt + enc.head(h_opt, batch["opt_mask"], h_opt, batch["opt_mask"]).to(h_opt.dtype) | |
| return h, h_opt.float() | |
| def encode_ctx(enc, batch, use_enc_head, train_enc): | |
| if train_enc: | |
| return encode(enc, batch, use_enc_head) | |
| with torch.no_grad(): | |
| return encode(enc, batch, use_enc_head) | |
| def readout_post(enc, hn): | |
| """S1 Yes/No readout on POST-norm states (pre-norm option states have norm ~1,700; post-norm ~46, so a head | |
| correction added post-norm has real leverage). With hn = norm(h_opt) this equals enc.readout(h_opt).""" | |
| w = enc.base.lm_head.weight | |
| rows = torch.stack([w[enc.yes_id], w[enc.no_id]]).float() | |
| z = F.linear(hn.float(), rows) | |
| return z[..., 0] - z[..., 1] | |
| def rollout(enc, R, batch, K, grad_steps, use_enc_head, reads=None, seed=None, train_enc=False): | |
| """Returns {k: slot scores} for k in reads (or the grad steps when training).""" | |
| h, h_opt = encode_ctx(enc, batch, use_enc_head, train_enc and torch.is_grad_enabled()) | |
| ctx = R.ctx_norm(h.float()) | |
| hn = enc.bb.norm(h_opt.to(h.dtype)).float() | |
| s = readout_post(enc, hn) | |
| outs = {0: s} if reads and 0 in reads else {} | |
| B = h.shape[0] | |
| g = torch.Generator(device=h.device) | |
| if seed is not None: | |
| g.manual_seed(seed) | |
| z = R.z0[None].expand(B, -1, -1) + 0.5 * torch.randn((B,) + R.z0.shape, device=h.device, generator=g if seed is not None else None) | |
| p = slot_beliefs(s, batch["items"]).detach() | |
| for k in range(1, K + 1): | |
| with torch.set_grad_enabled(torch.is_grad_enabled() and k > K - grad_steps): | |
| y = R.y_slots(h_opt, p) | |
| z = R.step(z, y, batch["opt_mask"], ctx, batch["attn"]) | |
| s = readout_post(enc, hn + R.delta(h_opt, z)) | |
| p = slot_beliefs(s, batch["items"]).detach() | |
| if (reads and k in reads) or (not reads and k > K - grad_steps): | |
| outs[k] = s | |
| return outs | |
| def item_loss(logits_list, items, device): | |
| tot = 0.0 | |
| for qs, x in zip(logits_list, items): | |
| per = [supervised_terms(lg.float(), torch.tensor([q["target"]], device=device), q["prim"], 0.25)["loss"] | |
| for lg, q in zip(qs, x["qs"])] | |
| tot = tot + sum(per) / len(per) | |
| return tot | |
| def load_encoder(path, use_head, device, mp, trainable=False): | |
| from eval import build_model | |
| enc = build_model(path, 0, use_head, device, mp, dtype=torch.float32 if trainable else torch.bfloat16) | |
| for p in enc.parameters(): | |
| p.requires_grad_(trainable) | |
| return enc.train() if trainable else enc.eval() | |
| def uses_head(cfg): | |
| return cfg.mode_data == "single" and cfg.encoder != "s1" | |
| def get_items(a, split, tok, mid, mt): | |
| if a.mode_data == "single": | |
| path = PROC + {"train": "/train.jsonl", "test": "/test.jsonl", "dev": "/dev.jsonl"}[split] | |
| if split == "train": | |
| return single_items(path, a.skip + a.n_items, a.skip, tok, mid, mt) | |
| return single_items(path, 0, 0, tok, mid, mt) | |
| raw = {"train": "train", "test": "test", "dev": "validation"}[split] | |
| canon = PROC + {"train": "/train.jsonl", "test": "/test.jsonl", "dev": "/dev.jsonl"}[split] | |
| if split == "train": | |
| return multi_items(raw, canon, a.skip, a.n_items, tok, mid, mt, True) | |
| return multi_items(raw, canon, 0, 0, tok, mid, mt, False) | |
| def batches(items, budget, max_items, shuffle_rng=None): | |
| order = sorted(range(len(items)), key=lambda i: len(items[i]["ids"])) | |
| out, cur = [], [] | |
| for i in order: | |
| n = len(items[i]["ids"]) | |
| if cur and (n * (len(cur) + 1) > budget or len(cur) >= max_items): | |
| out.append(cur) | |
| cur = [] | |
| cur.append(i) | |
| if cur: | |
| out.append(cur) | |
| if shuffle_rng is not None: | |
| shuffle_rng.shuffle(out) | |
| return out | |
| def train(a): | |
| device = torch.device("cuda") | |
| torch.manual_seed(a.seed) | |
| rng = random.Random(a.seed) | |
| tok, mp, mid, mt = tokenizer_and_mapping() | |
| items = get_items(a, "train", tok, mid, mt) | |
| use_head = uses_head(a) | |
| enc = load_encoder(a.encoder, use_head, device, mp, trainable=a.train_encoder) | |
| enc.compile_layers() | |
| R = Reasoner(n_lat=a.n_lat, n_layers=a.n_layers).to(device) | |
| ema = copy.deepcopy(R).eval() | |
| for p in ema.parameters(): | |
| p.requires_grad_(False) | |
| groups = [dict(params=list(R.parameters()), lr=a.lr, weight_decay=0.1, peak=a.lr)] | |
| if a.train_encoder: | |
| groups.append(dict(params=[p for p in enc.parameters() if p.requires_grad], lr=a.lr_enc, weight_decay=0.01, peak=a.lr_enc)) | |
| opt = torch.optim.AdamW(groups, fused=True) | |
| plan = batches(items, a.budget, a.batch, rng) | |
| total = len(plan) | |
| warm = max(1, int(0.05 * total)) | |
| print(json.dumps(dict(event="start", mode=a.mode_data, items=len(items), steps=total, | |
| questions=sum(len(x["qs"]) for x in items), params=sum(p.numel() for p in R.parameters()))), flush=True) | |
| log = open(os.path.join(a.out, "train_log.jsonl"), "a") | |
| t0 = time.time() | |
| for step, idx in enumerate(plan, 1): | |
| f = step / 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"] = g["peak"] * max(f, 0.02) | |
| mb = [items[i] for i in idx] | |
| b = collate(mb, device, tok.pad_token_id) | |
| K = rng.randint(1, a.k_max) | |
| gs = min(3, K) | |
| with torch.autocast("cuda", dtype=torch.bfloat16): | |
| outs = rollout(enc, R, b, K, gs, use_head, train_enc=a.train_encoder) | |
| losses = {k: item_loss(per_question(s.float(), mb), mb, device) / len(mb) for k, s in outs.items()} | |
| loss = losses[K] + (0.3 * sum(v for k, v in losses.items() if k != K) / max(1, len(losses) - 1) if len(losses) > 1 else 0.0) | |
| loss.backward() | |
| gn = torch.nn.utils.clip_grad_norm_([p for g in opt.param_groups for p in g["params"]], 1.0) | |
| opt.step() | |
| opt.zero_grad(set_to_none=True) | |
| with torch.no_grad(): | |
| for pe, pr in zip(ema.parameters(), R.parameters()): | |
| pe.mul_(a.ema).add_(pr.detach(), alpha=1 - a.ema) | |
| if step % 20 == 0 or step == total: | |
| rec = dict(step=step, total=total, K=K, loss_final=round(float(losses[K]), 4), grad_norm=round(float(gn), 4), | |
| head_o=round(R.head.o.weight.norm().item(), 3), elapsed=round(time.time() - t0, 1), | |
| 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=ema.state_dict(), cfg=vars(a)), os.path.join(a.out, "reasoner_ema.pt")) | |
| if a.train_encoder: | |
| from safetensors.torch import save_file | |
| sd = {k: v.detach().to(torch.bfloat16).contiguous() for k, v in enc.state_dict().items()} | |
| sd.pop("base.lm_head.weight", None) | |
| sd = {k: v for k, v in sd.items() if not k.startswith("bb.")} | |
| save_file(sd, os.path.join(a.out, "encoder.safetensors")) | |
| 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() | |
| ck = torch.load(os.path.join(a.out, "reasoner_ema.pt"), weights_only=False) | |
| cfg = argparse.Namespace(**ck["cfg"]) | |
| cfg.out = a.out | |
| items = get_items(cfg, a.split, tok, mid, mt) | |
| use_head = uses_head(cfg) | |
| enc_path = os.path.join(a.out, "encoder.safetensors") if getattr(cfg, "train_encoder", False) else cfg.encoder | |
| enc = load_encoder(enc_path, use_head, device, mp) | |
| enc.compile_layers() | |
| R = Reasoner(n_lat=cfg.n_lat, n_layers=cfg.n_layers).to(device) | |
| R.load_state_dict(ck["state"]) | |
| R.eval() | |
| reads = [0, 1, 2, 4, 8, 16, 32] | |
| probs = {} | |
| t0 = time.time() | |
| for idx in batches(items, 65536, 64): | |
| mb = [items[i] for i in idx] | |
| b = collate(mb, device, tok.pad_token_id) | |
| acc_runs = [] | |
| for seed in range(a.seeds): | |
| with torch.autocast("cuda", dtype=torch.bfloat16): | |
| outs = rollout(enc, R, b, max(reads), 0, use_head, reads=set(reads), seed=seed) | |
| acc_runs.append({k: per_question(s.float(), mb) for k, s in outs.items()}) | |
| for j, x in enumerate(mb): | |
| for qi, q in enumerate(x["qs"]): | |
| key = (x["id"], qi) | |
| probs[key] = {k: np.mean([torch.softmax(r[k][j][qi][0], -1).cpu().numpy() for r in acc_runs], 0) for k in reads} | |
| rows = [] | |
| for x in items: | |
| for qi, q in enumerate(x["qs"]): | |
| if q["kind"] != "hard": | |
| continue | |
| g = int(np.argmax(q["target"])) | |
| rows.append(dict(id=x["id"], qi=qi, task=x["task"], level=x["level"], is_target=q["is_target"], gold=g, | |
| n_q=len(x["qs"]), probs={str(k): [round(float(v), 6) for v in probs[(x["id"], qi)][k]] for k in reads})) | |
| with open(os.path.join(a.out, "pred_%s.jsonl" % a.split), "w") as f: | |
| for r in rows: | |
| f.write(json.dumps(r) + "\n") | |
| acc = lambda rs, k: 100 * float(np.mean([np.argmax(r["probs"][str(k)]) == r["gold"] for r in rs])) if rs else float("nan") | |
| rep = dict(split=a.split, seconds=round(time.time() - t0, 1), by_K={}) | |
| tgt = [r for r in rows if r["is_target"]] | |
| for k in reads: | |
| rep["by_K"][str(k)] = dict(all=acc(rows, k), target=acc(tgt, k), | |
| by_level={str(l): acc([r for r in rows if r["level"] == l], k) for l in range(5)}) | |
| print("K=%2d all %.2f target %.2f | %s" % (k, acc(rows, k), acc(tgt, k), | |
| " ".join("L%d %.2f" % (l, acc([r for r in rows if r["level"] == l], k)) for l in range(5))), flush=True) | |
| 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("--mode-data", choices=["single", "multi"], default="single") | |
| ap.add_argument("--encoder", default="") | |
| ap.add_argument("--skip", type=int, default=32000) | |
| ap.add_argument("--n-items", type=int, default=32000) | |
| ap.add_argument("--batch", type=int, default=64) | |
| ap.add_argument("--budget", type=int, default=65536) | |
| ap.add_argument("--lr", type=float, default=2e-4) | |
| ap.add_argument("--ema", type=float, default=0.995) | |
| ap.add_argument("--n-lat", type=int, default=32) | |
| ap.add_argument("--n-layers", type=int, default=2) | |
| ap.add_argument("--k-max", type=int, default=12) | |
| ap.add_argument("--seed", type=int, default=42) | |
| ap.add_argument("--split", choices=["test", "dev"], default="test") | |
| ap.add_argument("--seeds", type=int, default=4) | |
| ap.add_argument("--train-encoder", action="store_true") | |
| ap.add_argument("--lr-enc", type=float, default=2.5e-5) | |
| a = ap.parse_args() | |
| if not a.encoder: | |
| a.encoder = "/workspace/think-decide/runs/%s/model.safetensors" % ("p1" if a.mode_data == "single" else "mq") | |
| os.makedirs(a.out, exist_ok=True) | |
| train(a) if a.mode == "train" else evaluate(a) | |
Xet Storage Details
- Size:
- 17.6 kB
- Xet hash:
- d40e9492e180a2a97ba6d5eaee584b233706fbbe78dea76ed6257479c1a18b9e
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.