From 1460cf241f7c89daaf2f832098d0018b8c8e8dc6 Mon Sep 17 00:00:00 2001 From: Nils Date: Tue, 14 Jul 2026 11:22:46 +0200 Subject: [PATCH] GSM plan-distillation control (width/depth law falsification test) Co-Authored-By: Claude Fable 5 --- scripts/train_distill_gsm.py | 155 +++++++++++++++++++++++++++++++++++ 1 file changed, 155 insertions(+) create mode 100644 scripts/train_distill_gsm.py diff --git a/scripts/train_distill_gsm.py b/scripts/train_distill_gsm.py new file mode 100644 index 0000000..7dfd1dd --- /dev/null +++ b/scripts/train_distill_gsm.py @@ -0,0 +1,155 @@ +"""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"] + "", + 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()