156 lines
6.0 KiB
Python
156 lines
6.0 KiB
Python
"""GSM plan-distillation control: KL from CoT-context teacher into FF adapter.
|
|
|
|
The width/depth law's falsification test: on code, distilling the model's own
|
|
plans into a feedforward adapter beat every recurrent variant. If the law is
|
|
right, the same recipe FAILS on GSM — sequential arithmetic cannot be
|
|
compressed into weight-space plan-priming. Interpretation caveat (recorded):
|
|
GSM answer spans are ~3 tokens vs ~100 for code, so supervision density is
|
|
also lower here; a null is law-consistent but not law-proving.
|
|
"""
|
|
|
|
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 (COT_SUFFIX, DIRECT_SUFFIX, BandLooper, MergeAdapter,
|
|
chat_prompt, last_number, num_eq)
|
|
from prep_mbpp import batch_generate
|
|
|
|
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
|
|
|
|
|
|
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 ensure_cots(model, tok, items):
|
|
path = OUT / "gsm_cots.json"
|
|
if path.exists():
|
|
return json.load(open(path))
|
|
print(f"generating CoTs for {len(items)} items", flush=True)
|
|
gens = batch_generate(model, tok,
|
|
[chat_prompt(tok, it["question"], COT_SUFFIX)
|
|
for it in items], max_new_tokens=320,
|
|
batch_size=32)
|
|
cots = {}
|
|
for it, g in zip(items, gens):
|
|
if num_eq(last_number(g), it["gold"]): # keep only correct traces
|
|
cots[str(it["idx"])] = g
|
|
json.dump(cots, open(path, "w"), indent=1)
|
|
return cots
|
|
|
|
|
|
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 = [it for it in json.load(open(OUT / "star_data.json"))
|
|
if it["split"] == "train" and it["label"] != "drop"]
|
|
cots = ensure_cots(model, tok, data)
|
|
tok.padding_side = "right"
|
|
train = [it for it in data if str(it["idx"]) in cots
|
|
and len(tok(it["question"])["input_ids"]) <= 350]
|
|
print(f"distill pool: {len(train)}", flush=True)
|
|
|
|
pad = tok.pad_token_id or 0
|
|
t0 = time.time()
|
|
for step in range(STEPS):
|
|
batch = rng.sample(train, BATCH)
|
|
pairs = []
|
|
for it in batch:
|
|
a = tok(it["gold"] + "<end_of_turn>",
|
|
add_special_tokens=False)["input_ids"]
|
|
sp = tok(chat_prompt(tok, it["question"], DIRECT_SUFFIX),
|
|
add_special_tokens=False)["input_ids"]
|
|
tp = tok(chat_prompt(tok, it["question"], COT_SUFFIX)
|
|
+ cots[str(it["idx"])] + "\n",
|
|
add_special_tokens=False)["input_ids"]
|
|
pairs.append((sp, tp, a))
|
|
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):
|
|
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()
|
|
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 % 200 == 199 or step == STEPS - 1:
|
|
torch.save(adapter.state_dict(),
|
|
OUT / f"adapter_gsm_distill_e{step+1}.pt")
|
|
print("done", flush=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|