108 lines
4.1 KiB
Python
108 lines
4.1 KiB
Python
"""E1b (item 19): label-supervised halting head on a FROZEN trained merge.
|
|
|
|
The merge adapter (curriculum ★, adapter_code.pt) is loaded and frozen;
|
|
only the halting head trains. Targets from STaR labels: easy -> halt at
|
|
iteration 1, hard -> halt at iteration 4 (BCE at every iteration: 0 below
|
|
the target depth, 1 at/above). Mixed batches, no CE, no penalty — the
|
|
head is a depth-aware difficulty classifier on the loop trajectory.
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import random
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
from halting_common import HaltingMergeAdapter
|
|
from loop_common import BandLooper
|
|
from prep_mbpp import DIRECT_SUFFIX, mbpp_prompt
|
|
from train_merge_code import build_code_batch, MAX_TOK
|
|
|
|
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"))
|
|
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--merge", default=None,
|
|
help="frozen merge checkpoint (default results-loop/adapter_code.pt)")
|
|
ap.add_argument("--steps", type=int, default=300)
|
|
ap.add_argument("--kmax", type=int, default=4)
|
|
ap.add_argument("--seed", type=int, default=0)
|
|
ap.add_argument("--lr", type=float, default=3e-3)
|
|
ARGS = ap.parse_args()
|
|
BATCH = 4
|
|
TARGET_K = {"easy": 1, "hard": ARGS.kmax}
|
|
|
|
|
|
def main():
|
|
rng = random.Random(ARGS.seed)
|
|
torch.manual_seed(ARGS.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)
|
|
adapter = HaltingMergeAdapter(
|
|
d=model.config.get_text_config().hidden_size).cuda()
|
|
merge_ckpt = ARGS.merge or (OUT / "adapter_code.pt")
|
|
sd = torch.load(merge_ckpt, map_location="cuda")
|
|
adapter.load_state_dict(sd, strict=False) # merge weights; head stays init
|
|
for n, p in adapter.named_parameters():
|
|
p.requires_grad_(n.startswith("halt"))
|
|
opt = torch.optim.AdamW([p for p in adapter.parameters()
|
|
if p.requires_grad], lr=ARGS.lr)
|
|
print("trainable:", sum(p.numel() for p in adapter.parameters()
|
|
if p.requires_grad), flush=True)
|
|
|
|
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]
|
|
print(f"pool={len(train)}", flush=True)
|
|
|
|
log = []
|
|
t0 = time.time()
|
|
for step in range(ARGS.steps):
|
|
batch = rng.sample(train, BATCH)
|
|
ids, msk, lab, lmask = build_code_batch(tok, batch)
|
|
pl = (lmask.long().cumsum(-1).argmax(-1))
|
|
bidx = torch.arange(len(batch), device="cuda")
|
|
tk = torch.tensor([TARGET_K[it["label"]] for it in batch],
|
|
device="cuda")
|
|
with torch.no_grad():
|
|
calls, _ = looper.capture(ids, msk, logits_to_keep=1)
|
|
e = looper._hin[looper.l0].detach()
|
|
s = looper.band(e, calls)
|
|
loss = 0.0
|
|
for i in range(ARGS.kmax):
|
|
with torch.no_grad():
|
|
x = adapter(e, s)
|
|
x = torch.where(lmask[..., None], x, e)
|
|
s = looper.band(x, calls)
|
|
p = adapter.halt_prob(e[bidx, pl], s[bidx, pl])
|
|
tgt = ((i + 1) >= tk).float()
|
|
loss = loss + F.binary_cross_entropy(p.float(), tgt)
|
|
loss = loss / ARGS.kmax
|
|
opt.zero_grad(set_to_none=True)
|
|
loss.backward()
|
|
opt.step()
|
|
log.append({"step": step, "bce": loss.item()})
|
|
if step % 20 == 0:
|
|
print(f"step {step:4d} bce={loss.item():.4f} "
|
|
f"({(time.time()-t0)/(step+1):.1f}s/step)", flush=True)
|
|
torch.save(adapter.state_dict(), OUT / f"adapter_gatehead_e{ARGS.steps}.pt")
|
|
json.dump(log, open(OUT / "train_gatehead_log.json", "w"))
|
|
print("done", flush=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|