252 lines
11 KiB
Python
252 lines
11 KiB
Python
"""Train the merge adapter on MBPP with prompt-only ("latent planning") looping.
|
|
|
|
Same recipe as train_merge.py, with two changes:
|
|
- loop_mask: the merge applies only to prompt positions; the code tokens are
|
|
teacher-forced through the plain band (they still attend to looped prompt
|
|
states each iteration) — no exposure-bias gap by construction.
|
|
- supervision: the model's OWN passing code from prep_mbpp.py (easy: direct
|
|
pass; hard: code extracted from the plan-first pass), CE on code tokens.
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import math
|
|
import random
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
from loop_common import (AdaptiveMergeAdapter, BandLooper, MergeAdapter,
|
|
NoisyMergeAdapter, ParcaeAdapter, PerDepthAdapter,
|
|
RecurrentAdapter, TiedAlphaAdapter)
|
|
from prep_mbpp import DIRECT_SUFFIX, mbpp_prompt
|
|
|
|
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
|
|
WARMUP = 20
|
|
MAX_TOK = 512
|
|
K_BUCKETS = [(1, ("easy",)), (2, ("easy", "hard")), (4, ("hard",))]
|
|
# --deepk 16 rescales to [(2, easy), (8, mixed), (16, hard)]
|
|
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--seed", type=int, default=0)
|
|
ap.add_argument("--pause", type=int, default=0,
|
|
help="pause-token control: p inert tokens after prompt, "
|
|
"feedforward adapter, no recurrence")
|
|
ap.add_argument("--feedforward", action="store_true",
|
|
help="apply adapter once, no recurrence (pause-FF control)")
|
|
ap.add_argument("--alpha", type=float, default=0.3,
|
|
help="merge weight (2B-tuned default 0.3; try 0.1-0.15 at 12B)")
|
|
ap.add_argument("--adaptive", action="store_true",
|
|
help="state-dependent alpha (AdaptiveMergeAdapter)")
|
|
ap.add_argument("--deepk", type=int, default=0,
|
|
help="scale curriculum depths by deepk/4 (e.g. 16 -> 2/8/16)")
|
|
ap.add_argument("--bptt", type=int, default=0,
|
|
help="truncated BPTT: grads only through last N iterations")
|
|
ap.add_argument("--lr", type=float, default=1e-3)
|
|
ap.add_argument("--warm", default=None, help="warm-start adapter checkpoint")
|
|
ap.add_argument("--rec", action="store_true",
|
|
help="Huginn-style regime: RecurrentAdapter (learned A/B, "
|
|
"noise h0) + log-uniform random depth 1..recmax, "
|
|
"truncated bptt (default 4)")
|
|
ap.add_argument("--recmax", type=int, default=16)
|
|
ap.add_argument("--parcae", action="store_true",
|
|
help="rec regime with rho(A)<1 by construction "
|
|
"(diag-negative-exp ZOH parameterization)")
|
|
ap.add_argument("--randk", action="store_true",
|
|
help="factorial cell: standard MergeAdapter but rec-style "
|
|
"log-uniform random depth (isolates the curriculum)")
|
|
ap.add_argument("--noises0", action="store_true",
|
|
help="factorial cell: standard merge + noise s0")
|
|
ap.add_argument("--hidden", type=int, default=512,
|
|
help="adapter MLP hidden width (capacity control)")
|
|
ap.add_argument("--tiedalpha", action="store_true",
|
|
help="merge with learned per-dim constant alpha, B tied "
|
|
"to (1-a); standard curriculum")
|
|
ap.add_argument("--perdepth", type=int, default=0,
|
|
help="PerDepthAdapter: one merge per iteration depth "
|
|
"(Bae-style relaxation), standard curriculum")
|
|
ARGS, _ = ap.parse_known_args()
|
|
if ARGS.parcae:
|
|
ARGS.rec = True
|
|
if (ARGS.rec or ARGS.randk) and not ARGS.bptt:
|
|
ARGS.bptt = 4
|
|
SEED = ARGS.seed
|
|
LR = ARGS.lr
|
|
SUFFIX = ((f"_s{SEED}" if SEED else "")
|
|
+ (f"_p{ARGS.pause}" if ARGS.pause else "")
|
|
+ (f"_a{ARGS.alpha}" if ARGS.alpha != 0.3 else "")
|
|
+ ("_ad" if ARGS.adaptive else "")
|
|
+ (f"_dk{ARGS.deepk}" if ARGS.deepk else "")
|
|
+ (f"_lr{ARGS.lr}" if ARGS.lr != 1e-3 else "")
|
|
+ ("_warm" if ARGS.warm else "")
|
|
+ (("_parcae" if ARGS.parcae else "_rec") + str(ARGS.recmax)
|
|
if ARGS.rec else "")
|
|
+ (f"_pd{ARGS.perdepth}" if ARGS.perdepth else "")
|
|
+ ("_ta" if ARGS.tiedalpha else "")
|
|
+ (f"_rk{ARGS.recmax}" if ARGS.randk else "")
|
|
+ ("_ns" if ARGS.noises0 else "")
|
|
+ (f"_h{ARGS.hidden}" if ARGS.hidden != 512 else ""))
|
|
PAUSE_ID = 6 # <unused0>
|
|
|
|
|
|
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_code_batch(tok, items, device="cuda"):
|
|
"""Right-padded batch; labels on code tokens; loop_mask on prompt span."""
|
|
seqs, labs = [], []
|
|
for it in items:
|
|
p = tok(mbpp_prompt(tok, it, DIRECT_SUFFIX),
|
|
add_special_tokens=False)["input_ids"]
|
|
p = p + [PAUSE_ID] * ARGS.pause
|
|
a = tok("```python\n" + it["sol_code"] + "\n```<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))
|
|
|
|
|
|
@torch.no_grad()
|
|
def val_loss(looper, adapter, tok, items, k):
|
|
tot, n = 0.0, 0
|
|
for i in range(0, len(items), BATCH):
|
|
ids, msk, lab, lmask = build_code_batch(tok, items[i : i + BATCH])
|
|
logits = looper.loop_logits(adapter, ids, k, attention_mask=msk,
|
|
loop_mask=lmask, feedforward=ARGS.feedforward)
|
|
loss = F.cross_entropy(logits[:, :-1].flatten(0, 1).float(),
|
|
lab[:, 1:].flatten(), ignore_index=-100)
|
|
tot += loss.item() * len(ids)
|
|
n += len(ids)
|
|
return tot / n
|
|
|
|
|
|
def main():
|
|
rng = random.Random(SEED)
|
|
torch.manual_seed(SEED)
|
|
data = json.load(open(OUT / "mbpp_data.json"))
|
|
|
|
model, tok = load_model(dtype=torch.bfloat16)
|
|
for p in model.parameters():
|
|
p.requires_grad_(False)
|
|
looper = BandLooper(model)
|
|
d = model.config.get_text_config().hidden_size
|
|
H = ARGS.hidden
|
|
if ARGS.tiedalpha:
|
|
adapter = TiedAlphaAdapter(d=d, alpha=ARGS.alpha, hidden=H).cuda()
|
|
elif ARGS.noises0:
|
|
adapter = NoisyMergeAdapter(d=d, alpha=ARGS.alpha, hidden=H).cuda()
|
|
elif ARGS.perdepth:
|
|
adapter = PerDepthAdapter(d=d, alpha=ARGS.alpha,
|
|
n_depth=ARGS.perdepth).cuda()
|
|
elif ARGS.parcae:
|
|
adapter = ParcaeAdapter(d=d, alpha=ARGS.alpha).cuda()
|
|
elif ARGS.rec:
|
|
adapter = RecurrentAdapter(d=d, alpha=ARGS.alpha).cuda()
|
|
elif ARGS.adaptive:
|
|
adapter = AdaptiveMergeAdapter(d=d, alpha0=ARGS.alpha).cuda()
|
|
else:
|
|
adapter = MergeAdapter(d=d, alpha=ARGS.alpha, hidden=H).cuda()
|
|
if ARGS.warm:
|
|
adapter.load_state_dict(torch.load(ARGS.warm, map_location="cuda"))
|
|
print("warm-started from", ARGS.warm, flush=True)
|
|
global K_BUCKETS
|
|
if ARGS.deepk:
|
|
f = ARGS.deepk / 4
|
|
K_BUCKETS = [(max(1, int(k * f)), lbls) for k, lbls in K_BUCKETS]
|
|
print("K_BUCKETS ->", [(k, l) for k, l in K_BUCKETS], flush=True)
|
|
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["sol_code"]]
|
|
train = [it for it in train
|
|
if len(tok(mbpp_prompt(tok, it, DIRECT_SUFFIX))["input_ids"])
|
|
+ len(tok(it["sol_code"])["input_ids"]) + 12 <= MAX_TOK]
|
|
pool = {"easy": [it for it in train if it["label"] == "easy"],
|
|
"hard": [it for it in train if it["label"] == "hard"]}
|
|
val = {lbl: lst[:16] for lbl, lst in pool.items()}
|
|
for lbl in pool:
|
|
pool[lbl] = pool[lbl][16:]
|
|
print(f"train pool: easy={len(pool['easy'])} hard={len(pool['hard'])}",
|
|
flush=True)
|
|
|
|
log = []
|
|
t0 = time.time()
|
|
for step in range(STEPS):
|
|
if ARGS.rec or ARGS.randk:
|
|
# randomized depth, log-uniform in [1, recmax], any difficulty
|
|
k = min(ARGS.recmax,
|
|
max(1, int(math.exp(rng.uniform(0, math.log(ARGS.recmax))))))
|
|
labels = ("easy", "hard")
|
|
else:
|
|
k, labels = K_BUCKETS[step % len(K_BUCKETS)]
|
|
cand = [it for lbl in labels for it in pool[lbl]]
|
|
batch = rng.sample(cand, min(BATCH, len(cand)))
|
|
ids, msk, lab, lmask = build_code_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,
|
|
feedforward=ARGS.feedforward,
|
|
bptt=ARGS.bptt or None)
|
|
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:
|
|
rho = (f" rho(A)={adapter.rho():.3f}"
|
|
if hasattr(adapter, "rho") else "")
|
|
print(f"step {step:4d} k={k} loss={loss.item():.4f}{rho} "
|
|
f"({(time.time()-t0)/(step+1):.1f}s/step)", flush=True)
|
|
if rho:
|
|
log.append({"step": step, "rho": adapter.rho()})
|
|
if step % 100 == 99 or step == STEPS - 1:
|
|
vals = {}
|
|
for kk in ((0, 1, 2, 4, 8, 16) if ARGS.rec else (0, 1, 2, 4)):
|
|
vals[f"easy_k{kk}"] = val_loss(looper, adapter, tok,
|
|
val["easy"], kk)
|
|
vals[f"hard_k{kk}"] = val_loss(looper, adapter, tok,
|
|
val["hard"], kk)
|
|
print(f" val@{step}: " +
|
|
" ".join(f"{n}={v:.3f}" for n, v in vals.items()), flush=True)
|
|
log.append({"step": step, "val": vals})
|
|
torch.save(adapter.state_dict(),
|
|
OUT / f"adapter_code{SUFFIX}_e{step+1}.pt")
|
|
json.dump(log, open(OUT / f"train_code_log{SUFFIX}.json", "w"), indent=1)
|
|
|
|
torch.save(adapter.state_dict(), OUT / f"adapter_code{SUFFIX}.pt")
|
|
json.dump(log, open(OUT / f"train_code_log{SUFFIX}.json", "w"), indent=1)
|
|
print("done; adapter ->", OUT / f"adapter_code{SUFFIX}.pt")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|