137 lines
4.8 KiB
Python
137 lines
4.8 KiB
Python
"""Train the merge adapter on Blocksworld (prompt-only latent planning).
|
|
|
|
Same recipe as train_merge_code.py; supervision = the model's own verified
|
|
passing plan (direct for easy, CoT-derived for hard — the plan section only,
|
|
CE on plan tokens)."""
|
|
|
|
import argparse
|
|
import json
|
|
import math
|
|
import os
|
|
import random
|
|
import re
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
from bw_prep import DIRECT_SUFFIX, chat
|
|
from loop_common import BandLooper, MergeAdapter
|
|
|
|
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
|
|
K_BUCKETS = [(1, ("easy",)), (2, ("easy", "hard")), (4, ("hard",))]
|
|
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--seed", type=int, default=0)
|
|
ARGS = ap.parse_args()
|
|
MOVELIST_RE = re.compile(r"(?:^|\n)\s*1\s*[.)]", re.M)
|
|
|
|
|
|
def plan_only(text):
|
|
"""From a CoT output, keep from the numbered list onward."""
|
|
m = MOVELIST_RE.search(text)
|
|
return text[m.start():].strip() if m else text.strip()
|
|
|
|
|
|
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 build_batch(tok, items, device="cuda"):
|
|
seqs, labs = [], []
|
|
for it in items:
|
|
p = tok(chat(tok, it, DIRECT_SUFFIX),
|
|
add_special_tokens=False)["input_ids"]
|
|
a = tok(plan_only(it["sol_plan"]) + "<end_of_turn>",
|
|
add_special_tokens=False)["input_ids"]
|
|
seqs.append(p + a)
|
|
labs.append([-100] * len(p) + a)
|
|
T = max(len(s) for s in seqs)
|
|
pad = tok.pad_token_id or 0
|
|
ids = torch.full((len(seqs), T), pad, dtype=torch.long)
|
|
lab = torch.full((len(seqs), T), -100, dtype=torch.long)
|
|
msk = torch.zeros((len(seqs), T), dtype=torch.long)
|
|
for i, (s, l) in enumerate(zip(seqs, labs)):
|
|
ids[i, : len(s)] = torch.tensor(s)
|
|
lab[i, : len(s)] = torch.tensor(l)
|
|
msk[i, : len(s)] = 1
|
|
lmask = (lab == -100) & (msk == 1)
|
|
return ids.to(device), msk.to(device), lab.to(device), lmask.to(device)
|
|
|
|
|
|
def main():
|
|
rng = random.Random(ARGS.seed)
|
|
torch.manual_seed(ARGS.seed)
|
|
data = json.load(open(OUT / "bw_data.json"))
|
|
|
|
model, tok = load_model(dtype=torch.bfloat16)
|
|
for p in model.parameters():
|
|
p.requires_grad_(False)
|
|
looper = BandLooper(model)
|
|
adapter = MergeAdapter(
|
|
d=model.config.get_text_config().hidden_size).cuda()
|
|
opt = torch.optim.AdamW(adapter.parameters(), lr=LR, weight_decay=0.01)
|
|
|
|
train = [it for it in data if it["split"] == "train"
|
|
and it["label"] != "drop" and it.get("sol_plan")]
|
|
pool = {l: [it for it in train if it["label"] == l]
|
|
for l in ("easy", "hard")}
|
|
val = {l: pool[l][:12] for l in pool}
|
|
pool = {l: pool[l][12:] for l in pool}
|
|
print(f"pool: easy={len(pool['easy'])} hard={len(pool['hard'])}",
|
|
flush=True)
|
|
buckets = K_BUCKETS
|
|
if len(pool["easy"]) < BATCH:
|
|
# BW direct pass ~1%: easy bucket is empty by construction; train
|
|
# hard-only rather than aborting the arm
|
|
buckets = [(2, ("hard",)), (4, ("hard",))]
|
|
print(f"easy pool too small ({len(pool['easy'])}) — "
|
|
"hard-only curriculum k=2/4", flush=True)
|
|
if len(pool["hard"]) < BATCH:
|
|
print("INSUFFICIENT POOL — aborting", flush=True)
|
|
return
|
|
|
|
log = []
|
|
t0 = time.time()
|
|
for step in range(STEPS):
|
|
k, labels = buckets[step % len(buckets)]
|
|
cand = [it for lbl in labels for it in pool[lbl]]
|
|
batch = rng.sample(cand, BATCH)
|
|
ids, msk, lab, lmask = build_batch(tok, batch)
|
|
for g in opt.param_groups:
|
|
g["lr"] = lr_at(step)
|
|
logits = looper.loop_logits(adapter, ids, k, attention_mask=msk,
|
|
use_checkpoint=True, loop_mask=lmask)
|
|
loss = F.cross_entropy(logits[:, :-1].flatten(0, 1).float(),
|
|
lab[:, 1:].flatten(), ignore_index=-100)
|
|
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, "k": k, "loss": loss.item()})
|
|
if step % 10 == 0:
|
|
print(f"step {step:4d} k={k} loss={loss.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_bw_e{step+1}.pt")
|
|
json.dump(log, open(OUT / "train_bw_log.json", "w"), indent=1)
|
|
print("done", flush=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|