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

item 26 pre-registered: result-staging lens supervision at pre-'=' positions (no causal leakage — low loss requires computation); arms lg03 and lt03+lg03; item-25 in-flight note (lce 10.3→2.3)

Browse Source 
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Nils
2026-07-16 23:23:13 +02:00
co-authored by Claude Fable 5
parent 64abfe66f3
commit 0cffce876b
3 changed files with 132 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/jobs/zzz_r_rungb_lg.sh
+16
View File
@@ -0,0 +1,16 @@
# 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_lg*.json results-loop/eval_gsm_carrycot_b1_lt03_lg*.json results-loop/train_carrycot_b1_l*_log.json results-loop/adapter_carrycot_b1_l*_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
# arm a: staging-only
$P train_carry_cot.py --drop-steps 1 --warm-start $LOOP_OUT/adapter_carrycot_e400.pt \
--steps 200 --lr 3e-4 --lensteach-gen 0.3
$P eval_carry_cot.py --adapter $LOOP_OUT/adapter_carrycot_b1_lg03_e200.pt \
--tag gsm_carrycot_b1_lg03 --grid 0:0,2:12,2:16 --n 256
# arm b: full lens curriculum (pauses + staging)
$P train_carry_cot.py --drop-steps 1 --warm-start $LOOP_OUT/adapter_carrycot_e400.pt \
--steps 200 --lr 3e-4 --lensteach 0.3 --lensteach-gen 0.3
$P eval_carry_cot.py --adapter $LOOP_OUT/adapter_carrycot_b1_lt03_lg03_e200.pt \
--tag gsm_carrycot_b1_lt03_lg03 --grid 0:0,2:12,2:16 --n 256
scripts/train_carry_cot.py
+87 -15
View File
@@ -63,12 +63,19 @@ ap.add_argument("--lensteach", type=float, default=0.0, metavar="LAMBDA",
"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")
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"_ln{ARGS.lensnoise.replace(',', '_')}" if ARGS.lensnoise else "") + (
f"_s{ARGS.seed}" if ARGS.seed else "") + ARGS.tag_suffix
@@ -112,6 +119,41 @@ class LensNoiseWrapper(torch.nn.Module):
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
@@ -214,7 +256,7 @@ def main():
f"layers {lora_band_layers[0]}-{lora_band_layers[-1]}, "
f"lr={ARGS.lora_lr}", flush=True)
lens_teach = None
if ARGS.lensteach:
if ARGS.lensteach or ARGS.lensteach_gen:
from loop_common import BAND
J30 = torch.load(Path(__file__).resolve().parent.parent
/ "results/jbar.pt",
@@ -230,16 +272,28 @@ def main():
lg = softcap * torch.tanh(lg / softcap)
return lg
n_tgt = 0
for it in data:
if it["deleted"]:
it["lens_targets"] = tok(
it["deleted"], add_special_tokens=False
)["input_ids"][: it["extra_pauses"]]
n_tgt += bool(it["lens_targets"])
print(f"lens-teach λ={ARGS.lensteach}: targets on {n_tgt} items "
f"(pause j <-> deleted-step token j, lens at L{BAND[1]})",
flush=True)
if ARGS.lensteach:
n_tgt = 0
for it in data:
if it["deleted"]:
it["lens_targets"] = tok(
it["deleted"], add_special_tokens=False
)["input_ids"][: it["extra_pauses"]]
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)
groups = [{"params": list(adapter.parameters()), "lr": LR, "base": LR}]
if lora_params:
groups.append({"params": lora_params, "lr": ARGS.lora_lr,
@@ -270,7 +324,7 @@ def main():
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:
if ARGS.lensteach or ARGS.lensteach_gen:
logits, S = carry_logits(looper, adapter, ids, msk, plens,
K_PREFILL, use_checkpoint=True,
feedforward=ARGS.feedforward,
@@ -281,7 +335,7 @@ def main():
feedforward=ARGS.feedforward)
loss = F.cross_entropy(logits[:, :-1].flatten(0, 1).float(),
lab[:, 1:].flatten(), ignore_index=-100)
lce_val = 0.0
lce_val, lgen_val = 0.0, 0.0
if ARGS.lensteach:
terms = []
for b, it in enumerate(batch):
@@ -297,15 +351,33 @@ def main():
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
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})
log.append({"step": step, "loss": loss.item(), "lce": lce_val,
"lgen": lgen_val})
if step % 10 == 0:
print(f"step {step:4d} {lbl:4s} loss={loss.item():.4f} "
f"lce={lce_val:.3f} "
f"lce={lce_val:.3f} lgen={lgen_val:.3f} "
f"({(time.time()-t0)/(step+1):.1f}s/step)", flush=True)
if step % 200 == 199 or step == STEPS - 1:
if ARGS.lensnoise: