E1b trainer: label-supervised halting head on frozen merge (item 19)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

scripts/train_gate_head.py

@@ -0,0 +1,107 @@
+"""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()