jspace/scripts/train_distill.py

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