J-lens workspace reproduction + loop retrofit: lens, band looping, adapters, controls, multi-task evals
Reproduction of the 2026 workspace/J-lens paper on gemma-4 (E2B/12B/26B), plus the workspace-loop retrofit line: merge adapter, prompt-only latent planning (MBPP), carry variant, attribution controls (FF/pause/untrained), band-location ablation, Blocksworld harness, 12B replication scripts. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Nils
Claude Fable 5
@@ -0,0 +1,175 @@
|
||||
"""Plan-distillation baseline: same 1.6M budget, no recurrence.
|
||||
|
||||
Teacher: frozen model WITH its own terse plan in context (the plan text the
|
||||
STaR pipeline validated). Student: feedforward adapter (adapter(e,e) at
|
||||
prompt positions, applied once), NO plan in context. Loss: KL(teacher →
|
||||
student) on code tokens + 0.5 CE on gold code. Bounds how much of the loop's
|
||||
gain is available by compressing plan information into the adapter weights.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import random
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from loop_common import BandLooper, MergeAdapter
|
||||
from prep_mbpp import (DIRECT_SUFFIX, CODE_RE, batch_generate, mbpp_prompt)
|
||||
from prep_mbpp_fix import TERSE_PLAN_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
|
||||
KL_T = 1.0
|
||||
SEED = 0
|
||||
|
||||
|
||||
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 plan_prefix(gen_text):
|
||||
"""Text before the final fenced code block (the plan)."""
|
||||
m = list(CODE_RE.finditer(gen_text))
|
||||
return gen_text[: m[-1].start()].strip() if m else None
|
||||
|
||||
|
||||
def ensure_plans(model, tok, items):
|
||||
path = OUT / "mbpp_plans.json"
|
||||
if path.exists():
|
||||
plans = json.load(open(path))
|
||||
else:
|
||||
print(f"generating plans for {len(items)} items", flush=True)
|
||||
gens = batch_generate(model, tok,
|
||||
[mbpp_prompt(tok, it, TERSE_PLAN_SUFFIX)
|
||||
for it in items], max_new_tokens=700,
|
||||
batch_size=16)
|
||||
plans = {str(it["task_id"]): plan_prefix(g)
|
||||
for it, g in zip(items, gens)}
|
||||
json.dump(plans, open(path, "w"), indent=1)
|
||||
return plans
|
||||
|
||||
|
||||
def build_pair(tok, it, plan, device="cuda"):
|
||||
"""(student ids/mask/lmask, teacher ids, code positions in each)."""
|
||||
code = "```python\n" + it["sol_code"] + "\n```<end_of_turn>"
|
||||
a = tok(code, add_special_tokens=False)["input_ids"]
|
||||
sp = tok(mbpp_prompt(tok, it, DIRECT_SUFFIX),
|
||||
add_special_tokens=False)["input_ids"]
|
||||
tp = tok(mbpp_prompt(tok, it, TERSE_PLAN_SUFFIX) + plan + "\n",
|
||||
add_special_tokens=False)["input_ids"]
|
||||
return sp, tp, a
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--seed", type=int, default=0)
|
||||
args = ap.parse_args()
|
||||
rng = random.Random(args.seed)
|
||||
torch.manual_seed(args.seed)
|
||||
|
||||
model, tok = load_model(dtype=torch.bfloat16)
|
||||
tok.padding_side = "left"
|
||||
for p in model.parameters():
|
||||
p.requires_grad_(False)
|
||||
looper = BandLooper(model)
|
||||
d = model.config.get_text_config().hidden_size
|
||||
adapter = MergeAdapter(d=d).cuda()
|
||||
opt = torch.optim.AdamW(adapter.parameters(), lr=LR, weight_decay=0.01)
|
||||
|
||||
data = json.load(open(OUT / "mbpp_data.json"))
|
||||
train = [it for it in data if it["split"] == "train"
|
||||
and it["label"] != "drop" and it.get("sol_code")]
|
||||
plans = ensure_plans(model, tok, train)
|
||||
train = [it for it in train if plans.get(str(it["task_id"]))]
|
||||
tok.padding_side = "right"
|
||||
train = [it for it in train
|
||||
if len(tok(mbpp_prompt(tok, it, TERSE_PLAN_SUFFIX))["input_ids"])
|
||||
+ len(tok(plans[str(it["task_id"])])["input_ids"])
|
||||
+ len(tok(it["sol_code"])["input_ids"]) + 30 <= 900]
|
||||
print(f"distill pool: {len(train)}", flush=True)
|
||||
|
||||
pad = tok.pad_token_id or 0
|
||||
log = []
|
||||
t0 = time.time()
|
||||
for step in range(STEPS):
|
||||
batch = rng.sample(train, BATCH)
|
||||
pairs = [build_pair(tok, it, plans[str(it["task_id"])])
|
||||
for it in batch]
|
||||
# student tensors
|
||||
Ts = max(len(sp) + len(a) for sp, _, a in pairs)
|
||||
Tt = max(len(tp) + len(a) for _, tp, a in pairs)
|
||||
s_ids = torch.full((BATCH, Ts), pad, dtype=torch.long)
|
||||
s_msk = torch.zeros((BATCH, Ts), dtype=torch.long)
|
||||
t_ids = torch.full((BATCH, Tt), pad, dtype=torch.long)
|
||||
t_msk = torch.zeros((BATCH, Tt), dtype=torch.long)
|
||||
lab = torch.full((BATCH, Ts), -100, dtype=torch.long)
|
||||
spans = []
|
||||
for i, (sp, tp, a) in enumerate(pairs):
|
||||
s_ids[i, : len(sp) + len(a)] = torch.tensor(sp + a)
|
||||
s_msk[i, : len(sp) + len(a)] = 1
|
||||
lab[i, len(sp): len(sp) + len(a)] = torch.tensor(a)
|
||||
t_ids[i, : len(tp) + len(a)] = torch.tensor(tp + a)
|
||||
t_msk[i, : len(tp) + len(a)] = 1
|
||||
spans.append((len(sp), len(tp), len(a)))
|
||||
s_ids, s_msk, lab = s_ids.cuda(), s_msk.cuda(), lab.cuda()
|
||||
t_ids, t_msk = t_ids.cuda(), t_msk.cuda()
|
||||
lmask = (lab == -100) & (s_msk == 1)
|
||||
|
||||
with torch.no_grad():
|
||||
t_logits = model(input_ids=t_ids, attention_mask=t_msk,
|
||||
use_cache=False).logits
|
||||
for g in opt.param_groups:
|
||||
g["lr"] = lr_at(step)
|
||||
s_logits = looper.loop_logits(adapter, s_ids, 1, attention_mask=s_msk,
|
||||
loop_mask=lmask, feedforward=True,
|
||||
use_checkpoint=True)
|
||||
kl = torch.zeros((), device="cuda")
|
||||
n_tok = 0
|
||||
for i, (ls, lt, la) in enumerate(spans):
|
||||
# predicting code token j uses position (prefix_len + j - 1)
|
||||
sl = s_logits[i, ls - 1: ls + la - 1].float()
|
||||
tl = t_logits[i, lt - 1: lt + la - 1].float()
|
||||
kl = kl + F.kl_div(
|
||||
F.log_softmax(sl / KL_T, -1), F.log_softmax(tl / KL_T, -1),
|
||||
log_target=True, reduction="sum")
|
||||
n_tok += la
|
||||
kl = kl / n_tok
|
||||
ce = F.cross_entropy(s_logits[:, :-1].flatten(0, 1).float(),
|
||||
lab[:, 1:].flatten(), ignore_index=-100)
|
||||
loss = kl + 0.5 * ce
|
||||
opt.zero_grad(set_to_none=True)
|
||||
loss.backward()
|
||||
torch.nn.utils.clip_grad_norm_(adapter.parameters(), 1.0)
|
||||
opt.step()
|
||||
|
||||
log.append({"step": step, "kl": kl.item(), "ce": ce.item()})
|
||||
if step % 10 == 0:
|
||||
print(f"step {step:4d} kl={kl.item():.4f} ce={ce.item():.4f} "
|
||||
f"({(time.time()-t0)/(step+1):.1f}s/step)", flush=True)
|
||||
if step % 100 == 99 or step == STEPS - 1:
|
||||
torch.save(adapter.state_dict(),
|
||||
OUT / f"adapter_distill_e{step+1}.pt")
|
||||
json.dump(log, open(OUT / "train_distill_log.json", "w"),
|
||||
indent=1)
|
||||
|
||||
print("done", flush=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user