nils/jspace
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

item 27 pre-registered: zero-pause internal band looping — single M=10 in-place burst at prompt end, iteration-aligned lens targets, ii0 inference ablation; carry_steps gains inplace updates, trainer gains --inner-iters/--base-pauses

Browse Source 
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Nils
2026-07-17 00:21:20 +02:00
co-authored by Claude Fable 5
parent 0cffce876b
commit 44e814e78c
5 changed files with 149 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/carry_common.py
+48 -7
View File
@@ -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)
scripts/eval_carry_cot.py
+4 -1
View File
@@ -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)
scripts/jobs/zzz_s_rungb_ii.sh
+14
View File
@@ -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
scripts/train_carry_cot.py
+57 -8
View File
@@ -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