item 20 scored; item 21 (GSM carry-CoT + control) pre-registered and queued

scripts/eval_carry_cot.py

@@ -0,0 +1,94 @@
+"""E2 stage A eval (item 21): GSM8K accuracy for short-CoT-trained arms.
+
+Generates with generate_carry_c (prefill loop k + pause carry + per-token
+carry) or feedforward mode for the control arm; scores last_number vs
+gold, by STaR label. Grid over (k,p) cells.
+"""
+
+import argparse
+import json
+import os
+import sys
+import time
+from pathlib import Path
+
+import torch
+
+from carry_common import generate_carry_c
+from loop_common import (BandLooper, MergeAdapter, chat_prompt,
+                         DIRECT_SUFFIX, last_number, num_eq)
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+from jlens.core import load_model  # noqa: E402
+
+OUT = Path(os.environ.get("LOOP_OUT",
+    Path(__file__).resolve().parent.parent / "results-loop"))
+
+
+@torch.no_grad()
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--adapter", required=True)
+    ap.add_argument("--tag", required=True)
+    ap.add_argument("--grid", default="0:0,2:2,2:6",
+                    help="comma list of k:p cells")
+    ap.add_argument("--n", type=int, default=256)
+    ap.add_argument("--batch", type=int, default=8)
+    ap.add_argument("--feedforward", action="store_true")
+    ap.add_argument("--max-new", type=int, default=160)
+    args = ap.parse_args()
+
+    model, tok = load_model(dtype=torch.bfloat16)
+    tok.padding_side = "left"
+    looper = BandLooper(model)
+    adapter = MergeAdapter(
+        d=model.config.get_text_config().hidden_size).cuda()
+    adapter.load_state_dict(torch.load(args.adapter, map_location="cuda"))
+    adapter.eval()
+
+    items = [it for it in json.load(open(OUT / "star_data.json"))
+             if it["split"] == "test"][: args.n]
+    print(f"[{args.tag}] GSM carry-cot eval on {len(items)}, "
+          f"grid={args.grid} ff={args.feedforward}", flush=True)
+
+    res = {"tag": args.tag, "grid": {}, "n": len(items)}
+    for cell in args.grid.split(","):
+        k, p = (int(x) for x in cell.split(":"))
+        t0 = time.time()
+        hits, per_label, per_item = 0, {}, []
+        for i in range(0, len(items), args.batch):
+            chunk = items[i : i + args.batch]
+            enc = tok([chat_prompt(tok, it["question"], DIRECT_SUFFIX)
+                       for it in chunk], return_tensors="pt", padding=True,
+                      add_special_tokens=False).to("cuda")
+            if k == 0 and p == 0:
+                gen = model.generate(**enc, max_new_tokens=args.max_new,
+                                     do_sample=False)
+            else:
+                gen = generate_carry_c(looper, adapter, tok,
+                                       enc["input_ids"],
+                                       enc["attention_mask"], k, p,
+                                       max_new_tokens=args.max_new,
+                                       feedforward=args.feedforward)
+            for j, it in enumerate(chunk):
+                txt = tok.decode(gen[j, enc["input_ids"].shape[1]:],
+                                 skip_special_tokens=True)
+                ok = num_eq(last_number(txt), it["gold"])
+                hits += ok
+                d = per_label.setdefault(it["label"], [0, 0])
+                d[0] += ok
+                d[1] += 1
+                per_item.append({"idx": it["idx"], "ok": bool(ok)})
+        acc = hits / len(items)
+        by = {l: c / n for l, (c, n) in per_label.items()}
+        res["grid"][cell] = {"acc": acc, "by_label": by,
+                             "per_item": per_item}
+        print(f"{cell}: acc={acc:.3f}  "
+              f"by_label={ {l: round(v,3) for l,v in by.items()} }"
+              f"  ({time.time()-t0:.0f}s)", flush=True)
+    json.dump(res, open(OUT / f"eval_{args.tag}.json", "w"), indent=1)
+    print("wrote", OUT / f"eval_{args.tag}.json")
+
+
+if __name__ == "__main__":
+    main()

scripts/jobs/zzz_l_gsm_e2a.sh

@@ -0,0 +1,11 @@
+# gpuq-in:  results-loop/star_data.json
+# gpuq-out: results-loop/gsm_cot_data.json results-loop/eval_gsm_carrycot*.json results-loop/train_carrycot*_log.json
+git pull origin main -q 2>/dev/null
+P=/home/nils/jspace/.venv/bin/python
+export JLENS_MODEL=google/gemma-4-E2B-it LOOP_OUT=/home/nils/jspace/results-loop
+cd /home/nils/jspace/scripts
+$P prep_gsm_cot.py
+$P train_carry_cot.py
+$P eval_carry_cot.py --adapter $LOOP_OUT/adapter_carrycot_e400.pt --tag gsm_carrycot_e400 --grid 0:0,2:2,2:6 --n 256
+$P train_carry_cot.py --feedforward
+$P eval_carry_cot.py --adapter $LOOP_OUT/adapter_carrycot_ff_e400.pt --tag gsm_carrycot_ff_e400 --grid 2:2,2:6 --n 256 --feedforward

scripts/path_gates.py

@@ -0,0 +1,170 @@
+"""Learnable soft path through ALL layers (pre-registration item 14).
+
+The frozen model = 35 same-typed functions on one residual bus. Instead of
+a hand-fixed loop over L14-30, learn a gate matrix g[t, l] in [0,1]: on
+loop iteration t, layer l's residual delta is scaled by g[t, l] (prompt
+positions only; generated/answer positions always run ungated). Gates are
+initialized as a Gaussian bump over depth centered mid-band, so at init a
+loop iteration is approximately the hand band pass — then SGD may move the
+compute envelope anywhere in [0, n_layers). The anchor merge adapter is
+kept at each iteration boundary for stability (rho < 1).
+
+Reading the result: if the learned envelope concentrates on the lens band,
+gradient descent independently rediscovers the workspace; if it wins with
+mass elsewhere, the lens placement story needs revision.
+
+E2B caveat (stated in advance): KV sharing makes attention deltas of
+layers >= 15 loop-inert; gate mass there is interpretable for MLP deltas
+only.
+"""
+
+import math
+
+import torch
+import torch.nn as nn
+from torch.utils.checkpoint import checkpoint
+
+from loop_common import BAND, BandLooper, _text_model
+
+
+class PathGates(nn.Module):
+    """g[t, l] = sigmoid(logits[t, l]); row 0 = warm sweep, rows 1..k = loops."""
+
+    def __init__(self, n_layers, k_max, mu=None, sigma=6.0, band=BAND):
+        super().__init__()
+        mu = (band[0] + band[1]) / 2 if mu is None else mu
+        init = torch.empty(k_max + 1, n_layers)
+        for l in range(n_layers):
+            p = math.exp(-((l - mu) ** 2) / (2 * sigma ** 2))
+            p = min(max(p, 1e-3), 1 - 1e-3)
+            init[:, l] = math.log(p / (1 - p))
+        self.logits = nn.Parameter(init)
+
+    def g(self, t):
+        return torch.sigmoid(
+            self.logits[min(t, self.logits.shape[0] - 1)].float())
+
+    def envelope(self):
+        with torch.no_grad():
+            return torch.sigmoid(self.logits.float()).tolist()
+
+
+class GatedLooper(BandLooper):
+    """BandLooper over the FULL depth with per-iteration per-layer gates."""
+
+    def __init__(self, model, gates):
+        tm = _text_model(model)
+        super().__init__(model, band=(0, len(tm.layers) - 1))
+        self.gates = gates
+
+    def _gated(self, h, calls, t, loop_mask=None):
+        g = self.gates.g(t).to(h.device)
+        x = h
+        for i in range(self.l0, self.l1 + 1):
+            args, kwargs = calls[i]
+            out = self.tm.layers[i](x, *args, **kwargs)
+            if isinstance(out, tuple):
+                out = out[0]
+            gi = g[i].to(x.dtype)
+            if loop_mask is not None:  # ungated (g=1) off the prompt span
+                gi = torch.where(loop_mask[..., None], gi,
+                                 torch.ones_like(loop_mask[..., None],
+                                                 dtype=x.dtype))
+            x = x + gi * (out - x)
+        return x
+
+    def loop_logits(self, adapter, input_ids, k, attention_mask=None,
+                    use_checkpoint=False, return_states=False,
+                    last_only=False, loop_mask=None, feedforward=False,
+                    bptt=None):
+        calls, base_logits = self.capture(input_ids, attention_mask,
+                                          logits_to_keep=1 if last_only else 0)
+        if k == 0:
+            return (base_logits, None) if return_states else base_logits
+        del base_logits
+        e = self._hin[self.l0].detach()
+
+        def sweep(x, t):
+            if use_checkpoint:
+                return checkpoint(
+                    lambda x_: self._gated(x_, calls, t, loop_mask), x,
+                    use_reentrant=False)
+            return self._gated(x, calls, t, loop_mask)
+
+        s = sweep(e, 0)  # warm sweep, t=0 (gates trainable here too)
+        states = [s]
+        n_nograd = max(0, k - bptt) if bptt else 0
+        for i in range(k):
+            if i < n_nograd:
+                with torch.no_grad():
+                    x = adapter(e, s)
+                    if loop_mask is not None:
+                        x = torch.where(loop_mask[..., None], x, e)
+                    s = self._gated(x, calls, i + 1, loop_mask)
+                s = s.detach()
+                states.append(s)
+                continue
+            x = adapter(e, s)
+            if loop_mask is not None:
+                x = torch.where(loop_mask[..., None], x, e)
+            s = sweep(x, i + 1)
+            states.append(s)
+        logits = self.suffix_logits(s, calls, last_only=last_only)
+        return (logits, states) if return_states else logits
+
+    @torch.no_grad()
+    def generate_frozen_prompt(self, adapter, tok, input_ids, k,
+                               max_new_tokens=220, attention_mask=None,
+                               stop_strs=(), feedforward=False,
+                               conv_out=None):
+        """Per-layer KV write-in: run the final gated sweep recording every
+        layer's INPUT, then one native prefill with pre-forward hooks
+        swapping each layer's hidden_states to the recorded stream — the
+        cache then holds exactly the gated states; decode is native."""
+        if k == 0:
+            return super().generate_frozen_prompt(
+                adapter, tok, input_ids, 0, max_new_tokens=max_new_tokens,
+                attention_mask=attention_mask, stop_strs=stop_strs)
+        calls, _ = self.capture(input_ids, attention_mask, logits_to_keep=1)
+        e = self._hin[self.l0]
+        s = self._gated(e, calls, 0)
+        for i in range(k):
+            x = adapter(e, s)
+            if i < k - 1:
+                s = self._gated(x, calls, i + 1)
+        # final sweep: record per-layer inputs of the gated stream
+        xs = {}
+        g = self.gates.g(k).to(x.device)
+        h = x
+        for i in range(self.l0, self.l1 + 1):
+            xs[i] = h
+            args, kwargs = calls[i]
+            out = self.tm.layers[i](h, *args, **kwargs)
+            if isinstance(out, tuple):
+                out = out[0]
+            h = h + g[i].to(h.dtype) * (out - h)
+        del calls
+
+        P = input_ids.shape[1]
+        handles = []
+        for i in range(self.l0, self.l1 + 1):
+            def pre(mod, args, kwargs, i=i):
+                hh = kwargs.get("hidden_states",
+                                args[0] if args else None)
+                if hh is not None and hh.shape[1] == P:  # prefill only
+                    if "hidden_states" in kwargs:
+                        kwargs["hidden_states"] = xs[i].to(hh.dtype)
+                        return args, kwargs
+                    return (xs[i].to(hh.dtype),) + args[1:], kwargs
+                return None
+            handles.append(self.tm.layers[i].register_forward_pre_hook(
+                pre, with_kwargs=True))
+        try:
+            gen = self.model.generate(
+                input_ids=input_ids, attention_mask=attention_mask,
+                max_new_tokens=max_new_tokens, do_sample=False,
+                pad_token_id=tok.pad_token_id or 0)
+        finally:
+            for hd in handles:
+                hd.remove()
+        return gen

scripts/prep_gsm_cot.py

@@ -0,0 +1,66 @@
+"""E2 stage A prep (item 21): harvest TERSE verified CoTs for GSM8K train.
+
+For each non-drop train item, sample a compressed scratchpad ("at most 3
+short steps"), keep it only if the final number matches gold (STaR
+filter). Output: results-loop/gsm_cot_data.json rows
+{idx, label, cot} — the dense supervision the latent-mode arms never had.
+"""
+
+import json
+import os
+import sys
+import time
+from pathlib import Path
+
+import torch
+
+from loop_common import chat_prompt, last_number, num_eq
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+from jlens.core import load_model  # noqa: E402
+
+OUT = Path(os.environ.get("LOOP_OUT",
+    Path(__file__).resolve().parent.parent / "results-loop"))
+TERSE_SUFFIX = ("\nSolve in at most 3 short steps, one line each, digits "
+                "only (like '4*6=24'). Then give the last line exactly as "
+                "'Answer: N'.")
+BATCH = 16
+MAX_NEW = 120
+
+
+@torch.no_grad()
+def main():
+    model, tok = load_model(dtype=torch.bfloat16)
+    tok.padding_side = "left"
+    items = [it for it in json.load(open(OUT / "star_data.json"))
+             if it["split"] == "train" and it["label"] != "drop"]
+    print(f"harvesting terse CoTs for {len(items)} train items", flush=True)
+    rows, kept = [], 0
+    t0 = time.time()
+    for i in range(0, len(items), BATCH):
+        chunk = items[i : i + BATCH]
+        enc = tok([chat_prompt(tok, it["question"], TERSE_SUFFIX)
+                   for it in chunk], return_tensors="pt", padding=True,
+                  add_special_tokens=False).to("cuda")
+        gen = model.generate(**enc, max_new_tokens=MAX_NEW, do_sample=False)
+        for j, it in enumerate(chunk):
+            txt = tok.decode(gen[j, enc["input_ids"].shape[1]:],
+                             skip_special_tokens=True).strip()
+            ok = num_eq(last_number(txt), it["gold"])
+            if ok:
+                kept += 1
+                rows.append({"idx": it["idx"], "label": it["label"],
+                             "cot": txt})
+        if i % 80 == 0:
+            print(f"[{i+len(chunk)}/{len(items)}] kept={kept} "
+                  f"({time.time()-t0:.0f}s)", flush=True)
+    json.dump(rows, open(OUT / "gsm_cot_data.json", "w"), indent=1)
+    by = {}
+    for r in rows:
+        by[r["label"]] = by.get(r["label"], 0) + 1
+    print(f"wrote {len(rows)} verified terse CoTs {by} -> gsm_cot_data.json",
+          flush=True)
+
+
+if __name__ == "__main__":
+    main()

scripts/train_carry_cot.py

@@ -0,0 +1,159 @@
+"""E2 stage A (item 21): dense short-CoT supervision through the carry
+whiteboard on GSM8K.
+
+train_carry.py skeleton, one change that matters: CE targets are the
+model's own VERIFIED terse scratchpad + answer (gsm_cot_data.json,
+~30-60 tokens) instead of the ~3-token bare answer — the dense-output
+ingredient the latent GSM arms structurally lacked.
+Control: --feedforward = same supervision, adapter(e,e), no recurrence.
+"""
+
+import argparse
+import json
+import os
+import math
+import random
+import sys
+import time
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+
+from carry_common import PAUSE_ID, carry_logits
+from loop_common import BandLooper, MergeAdapter, chat_prompt, DIRECT_SUFFIX
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+from jlens.core import load_model  # noqa: E402
+
+OUT = Path(os.environ.get("LOOP_OUT",
+    Path(__file__).resolve().parent.parent / "results-loop"))
+STEPS = 600
+BATCH = 4
+LR = 1e-3
+WARMUP = 20
+K_PREFILL = 2
+P_BY_LABEL = {"easy": 2, "hard": 6}
+
+ap = argparse.ArgumentParser()
+ap.add_argument("--feedforward", action="store_true")
+ap.add_argument("--seed", type=int, default=0)
+ARGS = ap.parse_args()
+TAG = ("carrycot_ff" if ARGS.feedforward else "carrycot") + (
+    f"_s{ARGS.seed}" if ARGS.seed else "")
+
+
+def lr_at(step):
+    if step < WARMUP:
+        return LR * (step + 1) / WARMUP
+    t = (step - WARMUP) / max(1, STEPS - WARMUP)
+    return 1e-4 + 0.5 * (LR - 1e-4) * (1 + math.cos(math.pi * t))
+
+
+def build_batch(tok, items, p, device="cuda"):
+    seqs, labs, plens = [], [], []
+    for it in items:
+        pr = tok(chat_prompt(tok, it["question"], DIRECT_SUFFIX),
+                 add_special_tokens=False)["input_ids"]
+        a = tok(it["cot"] + "<end_of_turn>",
+                add_special_tokens=False)["input_ids"]
+        seqs.append(pr + [PAUSE_ID] * p + a)
+        labs.append([-100] * (len(pr) + p) + a)
+        plens.append(len(pr))
+    T = max(len(s) for s in seqs)
+    pad = tok.pad_token_id or 0
+    ids = torch.full((len(seqs), T), pad, dtype=torch.long)
+    lab = torch.full((len(seqs), T), -100, dtype=torch.long)
+    msk = torch.zeros((len(seqs), T), dtype=torch.long)
+    for i, (s, l) in enumerate(zip(seqs, labs)):
+        ids[i, : len(s)] = torch.tensor(s)
+        lab[i, : len(s)] = torch.tensor(l)
+        msk[i, : len(s)] = 1
+    return (ids.to(device), msk.to(device), lab.to(device),
+            torch.tensor(plens, device=device))
+
+
+@torch.no_grad()
+def val_loss(looper, adapter, tok, items, p):
+    tot, n = 0.0, 0
+    for i in range(0, len(items), BATCH):
+        ids, msk, lab, plens = build_batch(tok, items[i : i + BATCH], p)
+        logits = carry_logits(looper, adapter, ids, msk, plens, K_PREFILL,
+                              feedforward=ARGS.feedforward)
+        loss = F.cross_entropy(logits[:, :-1].flatten(0, 1).float(),
+                               lab[:, 1:].flatten(), ignore_index=-100)
+        tot += loss.item() * len(ids)
+        n += len(ids)
+    return tot / n
+
+
+def main():
+    rng = random.Random(ARGS.seed)
+    torch.manual_seed(ARGS.seed)
+    star = {it["idx"]: it for it in json.load(open(OUT / "star_data.json"))
+            if it["split"] == "train"}
+    cots = json.load(open(OUT / "gsm_cot_data.json"))
+    data = []
+    for r in cots:
+        it = star.get(r["idx"])
+        if it is None:
+            continue
+        data.append({"question": it["question"], "label": r["label"],
+                     "cot": r["cot"]})
+
+    model, tok = load_model(dtype=torch.bfloat16)
+    for pp in model.parameters():
+        pp.requires_grad_(False)
+    looper = BandLooper(model)
+    adapter = MergeAdapter(
+        d=model.config.get_text_config().hidden_size).cuda()
+    opt = torch.optim.AdamW(adapter.parameters(), lr=LR, weight_decay=0.01)
+
+    keep = [it for it in data
+            if len(tok(it["question"])["input_ids"])
+            + len(tok(it["cot"])["input_ids"]) + 16 <= 460]
+    rng.shuffle(keep)
+    pool = {l: [it for it in keep if it["label"] == l]
+            for l in ("easy", "hard")}
+    val = {l: pool[l][:12] for l in pool}
+    pool = {l: pool[l][12:] for l in pool}
+    print(f"pool: easy={len(pool['easy'])} hard={len(pool['hard'])}",
+          flush=True)
+    if min(len(v) for v in pool.values()) < BATCH:
+        print("INSUFFICIENT POOL — aborting", flush=True)
+        return
+
+    log = []
+    t0 = time.time()
+    for step in range(STEPS):
+        lbl = ("easy", "hard")[step % 2]
+        batch = rng.sample(pool[lbl], BATCH)
+        p = P_BY_LABEL[lbl]
+        ids, msk, lab, plens = build_batch(tok, batch, p)
+        for g in opt.param_groups:
+            g["lr"] = lr_at(step)
+        logits = carry_logits(looper, adapter, ids, msk, plens, K_PREFILL,
+                              use_checkpoint=True,
+                              feedforward=ARGS.feedforward)
+        loss = F.cross_entropy(logits[:, :-1].flatten(0, 1).float(),
+                               lab[:, 1:].flatten(), ignore_index=-100)
+        opt.zero_grad(set_to_none=True)
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(adapter.parameters(), 1.0)
+        opt.step()
+        log.append({"step": step, "loss": loss.item()})
+        if step % 10 == 0:
+            print(f"step {step:4d}  {lbl:4s} loss={loss.item():.4f}  "
+                  f"({(time.time()-t0)/(step+1):.1f}s/step)", flush=True)
+        if step % 200 == 199 or step == STEPS - 1:
+            for l in ("easy", "hard"):
+                v = val_loss(looper, adapter, tok, val[l], P_BY_LABEL[l])
+                print(f"  val@{step}: {l}={v:.3f}", flush=True)
+            torch.save(adapter.state_dict(),
+                       OUT / f"adapter_{TAG}_e{step+1}.pt")
+            json.dump(log, open(OUT / f"train_{TAG}_log.json", "w"))
+    print("done", flush=True)
+
+
+if __name__ == "__main__":
+    main()