diff --git a/results-loop/PROTOCOL_UNIFIED.md b/results-loop/PROTOCOL_UNIFIED.md index bd597db..4177d6d 100644 --- a/results-loop/PROTOCOL_UNIFIED.md +++ b/results-loop/PROTOCOL_UNIFIED.md @@ -732,3 +732,29 @@ auxiliary supervision, A2, E2-N. generation" in the free-running sense) is NOT this — that's a rollout-based follow-up (latent DAgger) if either lens arm moves. Job: scripts/jobs/zzz_r_rungb_lg.sh. + +27. **Internal band looping, zero pause tokens (pre-registered + 2026-07-17 ~01:10, before running; Nils's clarified structure, + option A confirmed via question: one silent burst before output).** + Architecture: NO pause tokens anywhere — sequence is prompt + + visible output only. After the k=2 prompt settle, the last prompt + position runs M=10 IN-PLACE band iterations (s^i seeds from the + position's own previous band output — carry_common inplace + updates), then generation proceeds with the normal single-pass + carry under the visible tokens. Vertical thought burst instead of + a horizontal pause tape: iterations leave no KV entries — only + the final state survives, a true internal loop (the C'-flavored + architecture the ladder never tested; every prior rung bought + compute with positions). Supervision: lens-CE iteration i <-> + deleted-step token i (λ=0.3, the stable value), same d=1 deletion, + warm-start rung-A e400, 200 steps, seed 0. Eval n=256: 0:0 + (sanity), 2:0 with --inner-iters 10 (matched), 2:0 with + --inner-iters 0 (ablation: does the burst matter at inference?). + References: positional d=1 31.6 (same ±5 / >36.6 bands) and the + ii0 ablation cell (isolates the burst's causal contribution). + Honest priors: item 25's 2:12 just landed at 31.2 (flat), so the + write-vs-compute dissociation likely carries over; the + architectural deltas that could matter here: no attention-tape + (forces state-borne computation rather than KV re-reading) and + anchor-at-prompt (iterations see the full settled question). + Job: scripts/jobs/zzz_s_rungb_ii.sh. diff --git a/scripts/carry_common.py b/scripts/carry_common.py index 0abc0c1..a102167 100644 --- a/scripts/carry_common.py +++ b/scripts/carry_common.py @@ -36,18 +36,26 @@ def prompt_prefill(looper, adapter, e, calls, prompt_mask, k): def carry_steps(looper, adapter, e, calls, S, X, step_updates, - use_checkpoint=False): - """Sequential scan. step_updates: list of (row_idx, pos) index tensors.""" - for rows, pos in step_updates: + use_checkpoint=False, iter_states=None): + """Sequential scan. step_updates: list of (row_idx, pos, [inplace]) + index tensors. inplace=True re-iterates the SAME position, seeding + from its own previous band output (internal band looping) instead of + its left neighbor. iter_states: optional list — every inplace + update's fresh band output at pos is appended (for lens losses).""" + for upd in step_updates: + rows, pos = upd[0], upd[1] + inplace = len(upd) > 2 and upd[2] if rows.numel() == 0: continue - seed = S[rows, pos - 1] + seed = S[rows, pos] if inplace else S[rows, pos - 1] x_new = adapter(e[rows, pos], seed) X = X.clone() X[rows, pos] = x_new.to(X.dtype) S = (checkpoint(lambda X_: looper.band(X_, calls), X, use_reentrant=False) if use_checkpoint else looper.band(X, calls)) + if inplace and iter_states is not None: + iter_states.append(S[rows, pos]) return S, X @@ -63,9 +71,28 @@ def build_step_updates(prompt_lens, total_lens, device): return updates +def splice_inner_iters(updates, inner_iters, inner_at, prompt_lens, dev, B): + """Insert inner_iters in-place band iterations at the (batch-uniform + offset) anchor position, right after the scan first settles it.""" + j_anchor = int((inner_at - prompt_lens).max()) + rows = torch.arange(B, device=dev) + inner = [(rows, inner_at.to(dev), True)] * inner_iters + if j_anchor < 0: + # anchor is the last PROMPT position (no pauses at all): + # iterate there before the scan enters the visible tokens + return inner + updates + out = [] + for j, u in enumerate(updates): + out.append(u) + if j == j_anchor: + out += inner + return out + + def carry_logits(looper, adapter, input_ids, attention_mask, prompt_lens, k, use_checkpoint=False, feedforward=False, - return_states=False): + return_states=False, inner_iters=0, inner_at=None, + iter_states=None): """Teacher-forced design-C forward (right-padded batch). feedforward=True: pause-token control — same positions get the adapter as @@ -87,15 +114,21 @@ def carry_logits(looper, adapter, input_ids, attention_mask, prompt_lens, S, X = prompt_prefill(looper, adapter, e, calls, prompt_mask, k) total_lens = attention_mask.sum(-1) updates = build_step_updates(prompt_lens.to(dev), total_lens.to(dev), dev) + if inner_iters and inner_at is not None: + updates = splice_inner_iters(updates, inner_iters, inner_at, + prompt_lens.to(dev), dev, + input_ids.shape[0]) S, X = carry_steps(looper, adapter, e, calls, S, X, updates, - use_checkpoint=use_checkpoint) + use_checkpoint=use_checkpoint, + iter_states=iter_states) out = looper.suffix_logits(S, calls) return (out, S) if return_states else out @torch.no_grad() def generate_carry_c(looper, adapter, tok, input_ids, attention_mask, - k, p, max_new_tokens=10, feedforward=False): + k, p, max_new_tokens=10, feedforward=False, + inner_iters=0): """Greedy design-C generation (left-padded batch, uniform positions). Appends p pause tokens, prefill-loops the prompt, carries through the @@ -126,6 +159,14 @@ def generate_carry_c(looper, adapter, tok, input_ids, attention_mask, updates = [(torch.arange(B, device=dev), torch.full((B,), n_prompt + j, device=dev, dtype=torch.long)) for j in range(p)] + if inner_iters: + anchor = torch.full((B,), n_prompt + p - 1, device=dev, + dtype=torch.long) + inner = [(torch.arange(B, device=dev), anchor, True) + for _ in range(inner_iters)] + # p=0: iterate at the last prompt position, before any + # visible token — no pause tokens involved + updates = (inner + updates) if p == 0 else (updates + inner) S, X = carry_steps(looper, adapter, e, calls, S, X, updates) else: X = torch.cat([X_store, e[:, X_store.shape[1]:]], 1) diff --git a/scripts/eval_carry_cot.py b/scripts/eval_carry_cot.py index ad074b4..9cdbfe3 100644 --- a/scripts/eval_carry_cot.py +++ b/scripts/eval_carry_cot.py @@ -38,6 +38,8 @@ def main(): ap.add_argument("--max-new", type=int, default=160) ap.add_argument("--bandlora", default=None, metavar="LORA_PT", help="load a lora_*_e*.pt loop-only band-LoRA checkpoint") + ap.add_argument("--inner-iters", type=int, default=0, metavar="M", + help="M in-place band iterations at the last pause") args = ap.parse_args() model, tok = load_model(dtype=torch.bfloat16) @@ -83,7 +85,8 @@ def main(): enc["input_ids"], enc["attention_mask"], k, p, max_new_tokens=args.max_new, - feedforward=args.feedforward) + feedforward=args.feedforward, + inner_iters=args.inner_iters) for j, it in enumerate(chunk): txt = tok.decode(gen[j, enc["input_ids"].shape[1]:], skip_special_tokens=True) diff --git a/scripts/jobs/zzz_s_rungb_ii.sh b/scripts/jobs/zzz_s_rungb_ii.sh new file mode 100644 index 0000000..2b49624 --- /dev/null +++ b/scripts/jobs/zzz_s_rungb_ii.sh @@ -0,0 +1,14 @@ +# gpuq-in: results-loop/star_data.json results-loop/gsm_cot_data.json results-loop/adapter_carrycot_e400.pt +# gpuq-out: results-loop/eval_gsm_carrycot_b1_ii*.json results-loop/train_carrycot_b1_lt03_ii10_p0_log.json results-loop/adapter_carrycot_b1_lt03_ii10_p0_e200.pt +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 train_carry_cot.py --drop-steps 1 --pause-per-step 0 --base-pauses 0 \ + --inner-iters 10 --lensteach 0.3 \ + --warm-start $LOOP_OUT/adapter_carrycot_e400.pt --steps 200 --lr 3e-4 +A=$LOOP_OUT/adapter_carrycot_b1_lt03_ii10_p0_e200.pt +$P eval_carry_cot.py --adapter $A --tag gsm_carrycot_b1_ii10 \ + --grid 0:0,2:0 --n 256 --inner-iters 10 +$P eval_carry_cot.py --adapter $A --tag gsm_carrycot_b1_ii10_ablate \ + --grid 2:0 --n 256 --inner-iters 0 diff --git a/scripts/train_carry_cot.py b/scripts/train_carry_cot.py index 2d575b7..3046451 100644 --- a/scripts/train_carry_cot.py +++ b/scripts/train_carry_cot.py @@ -68,6 +68,15 @@ ap.add_argument("--lensteach-gen", type=float, default=0.0, metavar="LAMBDA", "— 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)") ARGS = ap.parse_args() STEPS, LR = ARGS.steps, ARGS.lr TAG = ("carrycot_ff" if ARGS.feedforward else "carrycot") + ( @@ -76,6 +85,8 @@ TAG = ("carrycot_ff" if ARGS.feedforward else "carrycot") + ( 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"_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 @@ -189,9 +200,16 @@ def build_batch(tok, items, p, device="cuda"): 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) + 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) + 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) @@ -276,9 +294,10 @@ def main(): 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"][: it["extra_pauses"]] + )["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 " @@ -320,15 +339,27 @@ def main(): for step in range(STEPS): lbl = ("easy", "hard")[step % 2] batch = rng.sample(pool[lbl], BATCH) - p = P_BY_LABEL[lbl] + 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 - if ARGS.lensteach or ARGS.lensteach_gen: + 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: + 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) + return_states=True, **ckw) else: logits = carry_logits(looper, adapter, ids, msk, plens, K_PREFILL, use_checkpoint=True, @@ -336,7 +367,22 @@ def main(): 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: + 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") @@ -383,7 +429,10 @@ def main(): if ARGS.lensnoise: adapter.noise_on = False for l in ("easy", "hard"): - v = val_loss(looper, adapter, tok, val[l], P_BY_LABEL[l]) + 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