"""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```" 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()