"""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) ap.add_argument("--lensnoise", default=None, metavar="RANK,SCALE", help="E2-N1: inject noise into the carried state during " "training, shaped by the top-RANK sensitivity " "directions of jbar at the band entrance; noise norm " "= SCALE * per-position state norm (e.g. 32,0.05)") ap.add_argument("--drop-steps", type=int, default=0, metavar="D", help="E2-L rung B: delete the first D scratchpad steps, " "each replaced by --pause-per-step extra pauses") ap.add_argument("--pause-per-step", type=int, default=10, help="pauses per deleted step (median step = 10 tokens)") ap.add_argument("--warm-start", default=None, metavar="ADAPTER_PT") ap.add_argument("--steps", type=int, default=STEPS) ap.add_argument("--lr", type=float, default=LR) ap.add_argument("--tag-suffix", default="", help="appended to TAG (distinguish control variants)") ap.add_argument("--bandlora", type=int, default=0, metavar="RANK", help="item 24: loop-only LoRA (lora_band.LoopLoRA) on every " "band layer, uniform scale 1.0 — active only during " "band re-runs, k=0 stays bit-exact") ap.add_argument("--lora-lr", type=float, default=1e-3) ap.add_argument("--lensteach", type=float, default=0.0, metavar="LAMBDA", help="item 25: latent process supervision — lens-CE at the " "replacement pauses against the DELETED step's tokens " "(1:1 pause j <-> step token j), mixed at LAMBDA into " "the output CE") ap.add_argument("--lensteach-gen", type=float, default=0.0, metavar="LAMBDA", help="item 26: result-staging supervision during generation " "— at each visible scratchpad line's pre-'=' positions " "(result not yet in causal context), lens-CE the " "carried state against that line's result tokens") ap.add_argument("--inner-iters", type=int, default=0, metavar="M", help="item 27: M in-place band iterations at the anchor " "(internal looping, no extra tokens); with --lensteach " "the lens-CE aligns iteration i <-> deleted-step token " "i instead of pause positions") ap.add_argument("--base-pauses", type=int, default=-1, metavar="P", help="override P_BY_LABEL with a fixed pause count; 0 = NO " "pause tokens at all (inner iterations anchor on the " "last prompt position)") ap.add_argument("--teachstate", type=float, default=0.0, metavar="LAMBDA", help="item 28 (Nils's variant): teacher-state distillation " "— frozen warm-start adapter runs the FULL cot (step " "visible); its band-exit state at the deleted step's " "last token becomes the target; burst final iterate " "trained to it by cosine, weighted LAMBDA") ARGS = ap.parse_args() STEPS, LR = ARGS.steps, ARGS.lr TAG = ("carrycot_ff" if ARGS.feedforward else "carrycot") + ( f"_b{ARGS.drop_steps}" if ARGS.drop_steps else "") + ( f"_blr{ARGS.bandlora}" if ARGS.bandlora else "") + ( f"_lt{str(ARGS.lensteach).replace('.', '')}" if ARGS.lensteach else "") + ( f"_lg{str(ARGS.lensteach_gen).replace('.', '')}" if ARGS.lensteach_gen else "") + ( f"_ii{ARGS.inner_iters}" if ARGS.inner_iters else "") + ( f"_ts{str(ARGS.teachstate).replace('.', '')}" if ARGS.teachstate else "") + ( f"_p{ARGS.base_pauses}" if ARGS.base_pauses >= 0 else "") + ( f"_ln{ARGS.lensnoise.replace(',', '_')}" if ARGS.lensnoise else "") + ( f"_s{ARGS.seed}" if ARGS.seed else "") + ARGS.tag_suffix def drop_cot_steps(cot, d): """Delete the first d scratchpad lines (front-first: the deleted computation must ride the pause-chain before the visible remainder). Returns (new_cot, n_deleted, deleted_text); unparseable cots pass through intact.""" lines = [l for l in cot.split("\n") if l.strip()] ans = [l for l in lines if l.startswith("Answer")] steps = [l for l in lines if not l.startswith("Answer")] if len(ans) != 1 or not steps: return cot, 0, "" n = min(d, len(steps)) return "\n".join(steps[n:] + ans), n, "\n".join(steps[:n]) class LensNoiseWrapper(torch.nn.Module): """Perturb the carried state s (not the anchor e) before the merge, within the span of the lens's top-r readout-sensitive directions. Train-time only (noise_on flag); eval and checkpoints use .base.""" def __init__(self, base, jbar_layer, rank, scale): super().__init__() self.base = base self.scale = scale self.noise_on = True J = jbar_layer.float() _, _, Vt = torch.linalg.svd(J, full_matrices=False) self.register_buffer("V", Vt[:rank].T.contiguous()) # (d, r) def forward(self, e, s): if self.noise_on and self.scale > 0: z = torch.randn(*s.shape[:-1], self.V.shape[1], device=s.device, dtype=torch.float32) n = z @ self.V.T n = n * (s.float().norm(dim=-1, keepdim=True) * self.scale / (n.norm(dim=-1, keepdim=True) + 1e-6)) s = (s.float() + n).to(s.dtype) return self.base(e, s) def gen_staging_targets(tok, cot): """Item 26: per-line result-staging spans, computed in TOKEN space. For each cot line with '=', the pre-'=' positions (result not yet in causal context) target the line's result tokens; the Answer line's 'Answer:' positions target the answer tokens. Returns [(rel_positions, result_token_ids), ...] relative to tok(cot).""" ids = tok(cot, add_special_tokens=False)["input_ids"] decoded = [tok.decode([t]) for t in ids] out, line_start = [], 0 for i, d in enumerate(decoded + ["\n"]): if "\n" not in d and i < len(ids): continue line = list(range(line_start, min(i, len(ids)))) line_start = i + 1 if not line: continue text = "".join(decoded[j] for j in line) seps = [j for j in line if decoded[j].strip() == "="] if seps: sep = seps[-1] elif text.strip().startswith("Answer"): colons = [j for j in line if decoded[j].strip() == ":"] if not colons: continue sep = colons[-1] else: continue result = [ids[j] for j in line if j > sep and decoded[j].strip()] pre = [j for j in line if j <= sep] if result and pre: out.append((pre, result)) return out 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"] + "", add_special_tokens=False)["input_ids"] pi = p + it.get("extra_pauses", 0) seqs.append(pr + [PAUSE_ID] * pi + a) labs.append([-100] * (len(pr) + pi) + 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): chunk = items[i : i + BATCH] ids, msk, lab, plens = build_batch(tok, chunk, p) ckw = {} if ARGS.inner_iters: extras = torch.tensor([c.get("extra_pauses", 0) for c in chunk], device=plens.device) ckw = dict(inner_iters=ARGS.inner_iters, inner_at=plens + p + extras - 1) logits = carry_logits(looper, adapter, ids, msk, plens, K_PREFILL, feedforward=ARGS.feedforward, **ckw) 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 = [] n_dropped = 0 for r in cots: it = star.get(r["idx"]) if it is None: continue cot, ndel, deleted = (drop_cot_steps(r["cot"], ARGS.drop_steps) if ARGS.drop_steps else (r["cot"], 0, "")) n_dropped += ndel data.append({"question": it["question"], "label": r["label"], "cot": cot, "deleted": deleted, "extra_pauses": ndel * ARGS.pause_per_step}) if ARGS.drop_steps: print(f"rung B d={ARGS.drop_steps}: {n_dropped} steps deleted " f"across {len(data)} items " f"({ARGS.pause_per_step} pauses each)", flush=True) 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() if ARGS.lensnoise: r, sc = ARGS.lensnoise.split(",") jbar = torch.load(Path(__file__).resolve().parent.parent / "results/jbar.pt", map_location="cpu")["Jbar"] from loop_common import BAND adapter = LensNoiseWrapper(adapter, jbar[BAND[0]], int(r), float(sc)).cuda() print(f"lens-noise: rank={r} scale={sc} on jbar L{BAND[0]}", flush=True) if ARGS.warm_start: (adapter.base if ARGS.lensnoise else adapter).load_state_dict( torch.load(ARGS.warm_start, map_location="cuda")) print(f"warm-started from {ARGS.warm_start}", flush=True) lora_params, lora_band_layers = [], [] if ARGS.bandlora: from lora_band import inject_band_lora from loop_common import BAND lora_band_layers = list(range(BAND[0], BAND[1] + 1)) scales = {l: 1.0 for l in lora_band_layers} lora_params = inject_band_lora(looper.tm, BAND[0], scales, rank=ARGS.bandlora) for p in lora_params: p.data = p.data.cuda() print(f"band-lora r={ARGS.bandlora}: " f"{sum(p.numel() for p in lora_params)/1e6:.1f}M params, " f"layers {lora_band_layers[0]}-{lora_band_layers[-1]}, " f"lr={ARGS.lora_lr}", flush=True) lens_teach = None if ARGS.lensteach or ARGS.lensteach_gen: from loop_common import BAND J30 = torch.load(Path(__file__).resolve().parent.parent / "results/jbar.pt", map_location="cuda")["Jbar"][BAND[1]].float() tm = looper.tm softcap = model.config.get_text_config().final_logit_softcapping def lens_teach(h): proj = h.float() @ J30.T x = tm.norm(proj.to(tm.norm.weight.dtype)) lg = model.lm_head(x) if softcap: lg = softcap * torch.tanh(lg / softcap) return lg if ARGS.lensteach: n_tgt = 0 for it in data: if it["deleted"]: cap = ARGS.inner_iters or it["extra_pauses"] it["lens_targets"] = tok( it["deleted"], add_special_tokens=False )["input_ids"][:cap] n_tgt += bool(it["lens_targets"]) print(f"lens-teach λ={ARGS.lensteach}: targets on {n_tgt} " f"items (pause j <-> deleted-step token j, lens at " f"L{BAND[1]})", flush=True) if ARGS.lensteach_gen: n_gt, n_spans = 0, 0 for it in data: gt = gen_staging_targets(tok, it["cot"]) if gt: it["gen_targets"] = gt n_gt += 1 n_spans += len(gt) print(f"lens-teach-gen λ={ARGS.lensteach_gen}: staging targets " f"on {n_gt} items ({n_spans} line-spans; pre-'=' " f"positions target the line result)", flush=True) if ARGS.teachstate: assert ARGS.warm_start and ARGS.inner_iters, \ "teachstate needs --warm-start (frozen teacher) + --inner-iters" t0_ = time.time() todo = [it for it in data if it["deleted"]] pt = ARGS.base_pauses if ARGS.base_pauses >= 0 else 0 with torch.no_grad(): for i in range(0, len(todo), BATCH): chunk = todo[i:i + BATCH] full_items = [{"question": c["question"], "cot": c["deleted"] + "\n" + c["cot"], "extra_pauses": 0} for c in chunk] ids_, msk_, _, plens_ = build_batch(tok, full_items, pt) _, S_ = carry_logits(looper, adapter, ids_, msk_, plens_, K_PREFILL, return_states=True) for b, c in enumerate(chunk): nstep = len(tok(c["deleted"], add_special_tokens=False)["input_ids"]) pos = int(plens_[b]) + pt + nstep - 1 c["teacher_state"] = S_[b, pos].float().clone() print(f"teacher states: {len(todo)} captured ({time.time()-t0_:.0f}s;" f" frozen warm-start adapter, full cot, band-exit at the " f"deleted step's last token)", flush=True) groups = [{"params": list(adapter.parameters()), "lr": LR, "base": LR}] if lora_params: groups.append({"params": lora_params, "lr": ARGS.lora_lr, "base": ARGS.lora_lr}) opt = torch.optim.AdamW(groups, weight_decay=0.01) keep = [it for it in data if len(tok(it["question"])["input_ids"]) + len(tok(it["cot"])["input_ids"]) + it.get("extra_pauses", 0) + 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 = (ARGS.base_pauses if ARGS.base_pauses >= 0 else P_BY_LABEL[lbl]) ids, msk, lab, plens = build_batch(tok, batch, p) for g in opt.param_groups: g["lr"] = g["base"] * lr_at(step) / LR ckw, itstates = {}, None if ARGS.inner_iters: extras = torch.tensor([b_.get("extra_pauses", 0) for b_ in batch], device=plens.device) assert (extras == extras[0]).all(), \ "inner-iters needs a batch-uniform pause block" ckw = dict(inner_iters=ARGS.inner_iters, inner_at=plens + p + extras - 1) if ARGS.lensteach or ARGS.teachstate: itstates = [] ckw["iter_states"] = itstates if ARGS.lensteach or ARGS.lensteach_gen or ARGS.inner_iters: logits, S = carry_logits(looper, adapter, ids, msk, plens, K_PREFILL, use_checkpoint=True, feedforward=ARGS.feedforward, return_states=True, **ckw) else: 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) lce_val, lgen_val = 0.0, 0.0 if ARGS.lensteach and ARGS.inner_iters: terms = [] for b, it in enumerate(batch): tgt = it.get("lens_targets") if not tgt or not itstates: continue m = min(len(tgt), len(itstates)) hs = torch.stack([itstates[i][b] for i in range(m)]) terms.append(F.cross_entropy( lens_teach(hs).float(), torch.tensor(tgt[:m], device=hs.device))) if terms: lce = torch.stack(terms).mean() lce_val = lce.item() loss = loss + ARGS.lensteach * lce elif ARGS.lensteach: terms = [] for b, it in enumerate(batch): tgt = it.get("lens_targets") if not tgt: continue s0 = int(plens[b]) + p hs = S[b, s0 : s0 + len(tgt)] terms.append(F.cross_entropy( lens_teach(hs).float(), torch.tensor(tgt, device=hs.device))) if terms: lce = torch.stack(terms).mean() lce_val = lce.item() loss = loss + ARGS.lensteach * lce if ARGS.lensteach_gen: gterms = [] for b, it in enumerate(batch): base = int(plens[b]) + p + it.get("extra_pauses", 0) for pre, res in it.get("gen_targets", []): posl = torch.tensor([base + r for r in pre], device=S.device) lg = lens_teach(S[b, posl]).float() gterms.append(torch.stack([ F.cross_entropy( lg, torch.full((len(pre),), r, dtype=torch.long, device=S.device)) for r in res]).mean()) if gterms: gl = torch.stack(gterms).mean() lgen_val = gl.item() loss = loss + ARGS.lensteach_gen * gl lts_val = 0.0 if ARGS.teachstate and itstates: tterms = [] for b, it in enumerate(batch): T = it.get("teacher_state") if T is None: continue tterms.append(1 - F.cosine_similarity( itstates[-1][b].float(), T, dim=0)) if tterms: lt = torch.stack(tterms).mean() lts_val = lt.item() loss = loss + ARGS.teachstate * lt opt.zero_grad(set_to_none=True) loss.backward() torch.nn.utils.clip_grad_norm_( list(adapter.parameters()) + lora_params, 1.0) opt.step() log.append({"step": step, "loss": loss.item(), "lce": lce_val, "lgen": lgen_val, "lts": lts_val}) if step % 10 == 0: print(f"step {step:4d} {lbl:4s} loss={loss.item():.4f} " f"lce={lce_val:.3f} lgen={lgen_val:.3f} " f"lts={lts_val:.3f} " f"({(time.time()-t0)/(step+1):.1f}s/step)", flush=True) if step % 200 == 199 or step == STEPS - 1: if ARGS.lensnoise: adapter.noise_on = False for l in ("easy", "hard"): v = val_loss(looper, adapter, tok, val[l], ARGS.base_pauses if ARGS.base_pauses >= 0 else P_BY_LABEL[l]) print(f" val@{step}: {l}={v:.3f}", flush=True) if ARGS.lensnoise: adapter.noise_on = True sd = (adapter.base if ARGS.lensnoise else adapter).state_dict() torch.save(sd, OUT / f"adapter_{TAG}_e{step+1}.pt") if lora_params: torch.save({"rank": ARGS.bandlora, "band": lora_band_layers, "tensors": [p.detach().cpu() for p in lora_params]}, OUT / f"lora_{TAG}_e{step+1}.pt") json.dump(log, open(OUT / f"train_{TAG}_log.json", "w")) print("done", flush=True) if __name__ == "__main__": main()