"""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("--traj-tf", type=float, default=0.0, metavar="LAMBDA", help="item 29: teacher-forced TRANSITION loss — iteration i " "gets teacher waypoint T[i-1] as merge input, its band " "output is pulled onto T[i] (cosine); 10 independent " "supervised one-step regressions") ap.add_argument("--traj-fr", type=float, default=0.0, metavar="LAMBDA", help="item 29: free-running trajectory loss — the burst's " "own state s_i pulled onto waypoint T[i] (exposure-gap " "term)") ap.add_argument("--traj-span", choices=("step", "full"), default="step", help="waypoints sampled evenly across the deleted step " "(job 1, d=1) or the full CoT (job 2, answer-only)") ap.add_argument("--kvmem", type=int, default=0, metavar="CODE", help="item 31: KVMemoryAdapter (code width) — burst states " "become per-band-layer KV prefix entries readable by " "frozen attention during the answer scan") ap.add_argument("--freeze-merge", action="store_true", help="freeze the (warm-started) merge adapter; train only " "the kvmem adapter") ap.add_argument("--symchain", choices=("tf", "st"), default=None, help="item 32: discrete latent chain — each burst tick " "feeds back the lens-snapped token embedding through " "a zero-init projector. tf: teacher-forced symbols " "(ground-truth deleted-step tokens); st: free-running " "straight-through snaps") 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"_tj{ARGS.traj_span[0]}{str(ARGS.traj_tf).replace('.', '')}" f"_{str(ARGS.traj_fr).replace('.', '')}" if (ARGS.traj_tf or ARGS.traj_fr) else "") + ( f"_kvm{ARGS.kvmem}" if ARGS.kvmem else "") + ( f"_sc{ARGS.symchain}" if ARGS.symchain else "") + ( "_fm" if ARGS.freeze_merge 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 sym_burst_forward(looper, adapter, proj, ids, msk, plens, m, sym_tf, k, lens_fn, embed_w, start_id): """Item 32 forward: discrete latent chain at the anchor (teacher-forced or straight-through symbols), then the visible-token carry scan.""" from carry_common import (prompt_prefill, build_step_updates, carry_steps, sym_iterate) dev = ids.device calls, _ = looper.capture(ids, msk, logits_to_keep=1) e = looper._hin[looper.l0].detach() B = ids.shape[0] ar = torch.arange(ids.shape[1], device=dev) pmask = ar[None, :] < plens[:, None].to(dev) S, X = prompt_prefill(looper, adapter, e, calls, pmask, k) rows = torch.arange(B, device=dev) anchor = (plens - 1).to(dev) states = [] S, X = sym_iterate(looper, adapter, proj, e, calls, S, X, rows, anchor, m, lens_fn, embed_w, sym_tf=sym_tf, start_id=start_id, use_checkpoint=True, iter_states=states) total = msk.sum(-1) updates = build_step_updates(plens.to(dev), total.to(dev), dev) S, X = carry_steps(looper, adapter, e, calls, S, X, updates, use_checkpoint=True) return looper.suffix_logits(S, calls), states def traj_burst_forward(looper, adapter, ids, msk, plens, m, T, tf_on, k, mem_adapter=None): """Item 29 forward: optional teacher-forced transition predictions, then the free-running burst (whose final state seeds the answer scan, matching inference), then the visible-token carry. Returns (logits, tf_preds, fr_states).""" from torch.utils.checkpoint import checkpoint from carry_common import prompt_prefill, build_step_updates, carry_steps dev = ids.device calls, _ = looper.capture(ids, msk, logits_to_keep=1) e = looper._hin[looper.l0].detach() B = ids.shape[0] ar = torch.arange(ids.shape[1], device=dev) pmask = ar[None, :] < plens[:, None].to(dev) S, X = prompt_prefill(looper, adapter, e, calls, pmask, k) rows = torch.arange(B, device=dev) anchor = (plens - 1).to(dev) s0 = S[rows, anchor] def upd(S, X, seed): x_new = adapter(e[rows, anchor], seed) X = X.clone() X[rows, anchor] = x_new.to(X.dtype) S = checkpoint(lambda X_: looper.band(X_, calls), X, use_reentrant=False) return S, X tf_preds, fr_states = [], [] if tf_on: for i in range(m): seed = s0 if i == 0 else T[:, i - 1].to(s0.dtype) S, X = upd(S, X, seed) tf_preds.append(S[rows, anchor]) for i in range(m): seed = s0 if i == 0 else fr_states[-1] S, X = upd(S, X, seed) fr_states.append(S[rows, anchor]) if mem_adapter is not None: from kv_memory import arm_memory arm_memory(mem_adapter(torch.stack(fr_states, 1))) total = msk.sum(-1) updates = build_step_updates(plens.to(dev), total.to(dev), dev) S, X = carry_steps(looper, adapter, e, calls, S, X, updates, use_checkpoint=True) out = looper.suffix_logits(S, calls) # NOTE: memory stays armed through backward (checkpoint recompute # must see the same graph); the caller disarms after opt.step(). return out, tf_preds, fr_states 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) mem_adapter = None if ARGS.kvmem: from kv_memory import install, KVMemoryAdapter from loop_common import BAND install(model) mem_adapter = KVMemoryAdapter(model, band=BAND, code=ARGS.kvmem).cuda() print(f"kv-memory adapter: code={ARGS.kvmem}, " f"{sum(p.numel() for p in mem_adapter.parameters())/1e6:.1f}M " f"params, band {BAND}", flush=True) if ARGS.freeze_merge: for p_ in adapter.parameters(): p_.requires_grad_(False) print("merge adapter FROZEN", 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.traj_tf or ARGS.traj_fr: assert ARGS.warm_start and ARGS.inner_iters \ and ARGS.base_pauses == 0, \ "traj mode needs --warm-start, --inner-iters, --base-pauses 0" data = [it for it in data if it["deleted"]] m = ARGS.inner_iters t0_ = time.time() with torch.no_grad(): for i in range(0, len(data), BATCH): chunk = data[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, 0) _, S_ = carry_logits(looper, adapter, ids_, msk_, plens_, K_PREFILL, return_states=True) for b, c in enumerate(chunk): if ARGS.traj_span == "step": span = len(tok(c["deleted"], add_special_tokens=False)["input_ids"]) else: span = len(tok(c["deleted"] + "\n" + c["cot"], add_special_tokens=False)["input_ids"]) idxs = [round(j * (span - 1) / (m - 1)) for j in range(m)] pos = [int(plens_[b]) + jj for jj in idxs] c["traj_states"] = S_[b, pos].float().clone() print(f"trajectory waypoints: {len(data)} items x {m} states " f"({ARGS.traj_span} span, {time.time()-t0_:.0f}s; frozen " f"warm-start teacher, full cot)", flush=True) sym_proj, EMBED_W, START_ID = None, None, None if ARGS.symchain: assert ARGS.warm_start and ARGS.inner_iters and ARGS.lensteach, \ "symchain needs --warm-start, --inner-iters, --lensteach" d_ = model.config.get_text_config().hidden_size sym_proj = torch.nn.Linear(d_, d_).cuda() torch.nn.init.zeros_(sym_proj.weight) torch.nn.init.zeros_(sym_proj.bias) EMBED_W = model.get_input_embeddings().weight.detach() START_ID = tok("\n", add_special_tokens=False)["input_ids"][0] print(f"symchain [{ARGS.symchain}]: zero-init proj " f"({d_}x{d_}), start_id={START_ID}", 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 = ([] if ARGS.freeze_merge else [{"params": list(adapter.parameters()), "lr": LR, "base": LR}]) if mem_adapter is not None: groups.append({"params": list(mem_adapter.parameters()), "lr": LR, "base": LR}) if sym_proj is not None: groups.append({"params": list(sym_proj.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, tfp, frs = {}, None, None, None if ARGS.symchain: sym_tf = None if ARGS.symchain == "tf": mat = torch.full((len(batch), ARGS.inner_iters), START_ID, dtype=torch.long) for b_i, it in enumerate(batch): tgt = it.get("lens_targets") or [] row = [START_ID] + tgt[:-1] mat[b_i, :len(row)] = torch.tensor(row) sym_tf = mat.cuda() logits, itstates = sym_burst_forward( looper, adapter, sym_proj, ids, msk, plens, ARGS.inner_iters, sym_tf, K_PREFILL, lens_teach, EMBED_W, START_ID) elif ARGS.traj_tf or ARGS.traj_fr or ARGS.kvmem: T = (torch.stack([torch.as_tensor(b_["traj_states"]) for b_ in batch]).cuda() if (ARGS.traj_tf or ARGS.traj_fr) else None) logits, tfp, frs = traj_burst_forward( looper, adapter, ids, msk, plens, ARGS.inner_iters, T, bool(ARGS.traj_tf), K_PREFILL, mem_adapter=mem_adapter) else: 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 ltf_val, lfr_val = 0.0, 0.0 if ARGS.traj_tf and tfp: ltf = torch.stack([ (1 - F.cosine_similarity(p_.float(), T[:, i], dim=-1)).mean() for i, p_ in enumerate(tfp)]).mean() ltf_val = ltf.item() loss = loss + ARGS.traj_tf * ltf if ARGS.traj_fr and frs: lfr = torch.stack([ (1 - F.cosine_similarity(s_.float(), T[:, i], dim=-1)).mean() for i, s_ in enumerate(frs)]).mean() lfr_val = lfr.item() loss = loss + ARGS.traj_fr * lfr opt.zero_grad(set_to_none=True) loss.backward() torch.nn.utils.clip_grad_norm_( [p_ for p_ in adapter.parameters() if p_.requires_grad] + lora_params + (list(mem_adapter.parameters()) if mem_adapter is not None else []) + (list(sym_proj.parameters()) if sym_proj is not None else []), 1.0) opt.step() if ARGS.kvmem: from kv_memory import arm_memory arm_memory(None) log.append({"step": step, "loss": loss.item(), "lce": lce_val, "lgen": lgen_val, "lts": lts_val, "ltf": ltf_val, "lfr": lfr_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} ltf={ltf_val:.3f} " f"lfr={lfr_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 mem_adapter is not None: torch.save(mem_adapter.state_dict(), OUT / f"kvmem_{TAG}_e{step+1}.pt") if sym_proj is not None: torch.save(sym_proj.state_dict(), OUT / f"symproj_{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()