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
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Nils
Claude Fable 5
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user