From 8fe8e34b4777bf591c4b9529c7770483b2ce32af Mon Sep 17 00:00:00 2001 From: Nils Date: Tue, 14 Jul 2026 11:00:42 +0200 Subject: [PATCH] rung-2 pilot: loop-only entrance-faded band LoRA, warm-start, inline eval; BW hard-only fallback Co-Authored-By: Claude Fable 5 --- scripts/loop_common.py | 9 ++- scripts/lora_band.py | 68 ++++++++++++++++++ scripts/train_rung2.py | 159 +++++++++++++++++++++++++++++++++++++++++ 3 files changed, 235 insertions(+), 1 deletion(-) create mode 100644 scripts/lora_band.py create mode 100644 scripts/train_rung2.py diff --git a/scripts/loop_common.py b/scripts/loop_common.py index 294925a..fb46dc6 100644 --- a/scripts/loop_common.py +++ b/scripts/loop_common.py @@ -132,7 +132,14 @@ class BandLooper: return h def band(self, h, calls): - return self._run(h, calls, self.l0, self.l1) + try: + from lora_band import loop_active + loop_active(True) + out = self._run(h, calls, self.l0, self.l1) + loop_active(False) + return out + except ImportError: + return self._run(h, calls, self.l0, self.l1) def suffix_logits(self, h, calls, last_only=False): h = self._run(h, calls, self.l1 + 1, self.n_layers - 1) diff --git a/scripts/lora_band.py b/scripts/lora_band.py new file mode 100644 index 0000000..33b2456 --- /dev/null +++ b/scripts/lora_band.py @@ -0,0 +1,68 @@ +"""Loop-only LoRA for band layers: deltas active ONLY during band re-runs. + +The initial capture forward runs the pristine base model (k=0 stays +bit-exact); BandLooper.band() re-runs flip LOOP_ACTIVE on, so the LoRA only +shapes the recurrence. Entrance-weighted per the rung-2 design note. +""" + +import torch +import torch.nn as nn + +_STATE = {"on": False} + + +class LoopLoRA(nn.Module): + """Wraps an nn.Linear; adds scaled low-rank delta when loop is active.""" + + def __init__(self, base: nn.Linear, rank=8, scale=1.0): + super().__init__() + self.base = base + self.rank = rank + self.scale = scale + self.A = nn.Parameter( + torch.randn(rank, base.in_features, dtype=torch.float32) * 0.01) + self.B = nn.Parameter( + torch.zeros(base.out_features, rank, dtype=torch.float32)) + + def forward(self, x): + out = self.base(x) + if _STATE["on"]: + delta = (x.float() @ self.A.T) @ self.B.T + out = out + (self.scale * delta).to(out.dtype) + return out + + +def loop_active(on: bool): + _STATE["on"] = on + + +def inject_band_lora(tm, l0, layers_scales, rank=8, targets=("q_proj", "v_proj", + "down_proj")): + """Wrap target Linears in given band layers. layers_scales: {layer: scale}. + Returns list of LoRA params.""" + params = [] + for li, scale in layers_scales.items(): + layer = tm.layers[li] + for name, mod in list(layer.named_modules()): + leaf = name.split(".")[-1] + if leaf in targets and isinstance(mod, nn.Linear): + parent = layer + parts = name.split(".") + for p in parts[:-1]: + parent = getattr(parent, p) + wrapped = LoopLoRA(mod, rank=rank, scale=scale) + setattr(parent, parts[-1], wrapped) + params += [wrapped.A, wrapped.B] + return params + + +def entrance_faded_scales(l0, fade_to, full_until=None): + """scale 1.0 at l0, linear fade to 0 by fade_to (exclusive).""" + full_until = full_until if full_until is not None else l0 + out = {} + for li in range(l0, fade_to): + if li <= full_until: + out[li] = 1.0 + else: + out[li] = max(0.0, 1.0 - (li - full_until) / (fade_to - full_until)) + return {li: s for li, s in out.items() if s > 0} diff --git a/scripts/train_rung2.py b/scripts/train_rung2.py new file mode 100644 index 0000000..ef4f4e1 --- /dev/null +++ b/scripts/train_rung2.py @@ -0,0 +1,159 @@ +"""Rung-2 pilot: warm-start trained adapter, unfreeze band via loop-only LoRA. + +Entrance-weighted LoRA (rank 8, q/v/down of L14-22, scale fading 1->0) active +ONLY inside band re-runs — k=0 remains bit-exact base model. Joint training +of adapter (warm) + LoRA on the standard MBPP curriculum; inline slow-path +eval (LoRA applies to every band traversal, so the fast prefill trick would +be inconsistent). +""" + +import argparse +import json +import math +import os +import random +import sys +import time +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import torch +import torch.nn.functional as F + +from lora_band import entrance_faded_scales, inject_band_lora +from loop_common import BandLooper, MergeAdapter, BAND +from prep_mbpp import DIRECT_SUFFIX, extract_code, mbpp_prompt, run_tests +from train_merge_code import build_code_batch # noqa: F401 (shared format) + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) +from jlens.core import load_model, _text_model # noqa: E402 + +OUT = Path(os.environ.get("LOOP_OUT", + Path(__file__).resolve().parent.parent / "results-loop")) +STEPS = 600 +BATCH = 4 +WARMUP = 20 +K_BUCKETS = [(1, ("easy",)), (2, ("easy", "hard")), (4, ("hard",))] + +ap = argparse.ArgumentParser() +ap.add_argument("--warm-adapter", default=None, + help="path to trained MergeAdapter to warm-start from") +ap.add_argument("--lr", type=float, default=3e-4) +ap.add_argument("--rank", type=int, default=8) +ap.add_argument("--seed", type=int, default=0) +ap.add_argument("--eval-n", type=int, default=250) +ARGS = ap.parse_args() + + +def lr_at(step): + if step < WARMUP: + return ARGS.lr * (step + 1) / WARMUP + t = (step - WARMUP) / max(1, STEPS - WARMUP) + return 0.1 * ARGS.lr + 0.45 * ARGS.lr * (1 + math.cos(math.pi * t)) + + +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) + tm = _text_model(model) + looper = BandLooper(model) + d = model.config.get_text_config().hidden_size + adapter = MergeAdapter(d=d).cuda() + if ARGS.warm_adapter: + adapter.load_state_dict(torch.load(ARGS.warm_adapter, + map_location="cuda")) + print("warm-started adapter from", ARGS.warm_adapter, flush=True) + + scales = entrance_faded_scales(BAND[0], BAND[0] + 9, + full_until=BAND[0] + 3) + lora_params = inject_band_lora(tm, BAND[0], scales, rank=ARGS.rank) + model.cuda() + print(f"LoRA on layers {sorted(scales)} " + f"({sum(p.numel() for p in lora_params)/1e6:.1f}M params)", + flush=True) + opt = torch.optim.AdamW( + [{"params": adapter.parameters(), "lr": ARGS.lr}, + {"params": lora_params, "lr": ARGS.lr}], weight_decay=0.01) + + 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 <= 512] + pool = {"easy": [it for it in train if it["label"] == "easy"][16:], + "hard": [it for it in train if it["label"] == "hard"][16:]} + print(f"pool: easy={len(pool['easy'])} hard={len(pool['hard'])}", + flush=True) + + t0 = time.time() + for step in range(STEPS): + k, labels = K_BUCKETS[step % len(K_BUCKETS)] + cand = [it for lbl in labels for it in pool[lbl]] + batch = rng.sample(cand, min(BATCH, len(cand))) + ids, msk, lab, lmask = build_code_batch(tok, batch) + for g in opt.param_groups: + g["lr"] = lr_at(step) + logits = looper.loop_logits(adapter, ids, k, attention_mask=msk, + use_checkpoint=True, loop_mask=lmask) + loss = F.cross_entropy(logits[:, :-1].flatten(0, 1).float(), + lab[:, 1:].flatten(), ignore_index=-100) + if not torch.isfinite(loss): + print(f"NON-FINITE LOSS step {step} — halting (McLeish warning)", + flush=True) + break + opt.zero_grad(set_to_none=True) + loss.backward() + torch.nn.utils.clip_grad_norm_( + [p for g in opt.param_groups for p in g["params"]], 1.0) + opt.step() + if step % 10 == 0: + print(f"step {step:4d} k={k} loss={loss.item():.4f} " + f"({(time.time()-t0)/(step+1):.1f}s/step)", flush=True) + if step % 200 == 199 or step == STEPS - 1: + torch.save({"adapter": adapter.state_dict(), + "lora": [ (p.detach().cpu()) for p in lora_params ]}, + OUT / f"rung2_e{step+1}.pt") + + # inline slow-path eval (LoRA active in every band traversal) + items = [it for it in data if it["split"] == "test"][: ARGS.eval_n] + res = {} + for k in (0, 2, 4): + codes = [] + for i in range(0, len(items), 8): + torch.cuda.empty_cache() + chunk = items[i : i + 8] + enc = tok([mbpp_prompt(tok, it, DIRECT_SUFFIX) for it in chunk], + return_tensors="pt", padding=True, + add_special_tokens=False).to("cuda") + gen = looper.loop_generate(adapter, tok, enc["input_ids"], k, + max_new_tokens=220, + attention_mask=enc["attention_mask"], + loop_prompt_only=True) + for j in range(len(chunk)): + codes.append(extract_code( + tok.decode(gen[j, enc["input_ids"].shape[1]:], + skip_special_tokens=True))) + with ThreadPoolExecutor(8) as ex: + oks = list(ex.map(lambda ci: run_tests(ci[0], ci[1]), + zip(codes, items))) + by = {} + for it, ok in zip(items, oks): + dd = by.setdefault(it["label"], [0, 0]) + dd[0] += ok + dd[1] += 1 + res[k] = {"acc": sum(oks) / len(items), + "by_label": {l: c / n for l, (c, n) in by.items()}} + print(f"k={k}: pass@1={res[k]['acc']:.3f} " + f"by_label={ {l: round(v,3) for l,v in res[k]['by_label'].items()} }", + flush=True) + json.dump(res, open(OUT / "eval_rung2.json", "w"), indent=1) + print("wrote", OUT / "eval_rung2.json") + + +if __name__ == "__main__": + main()