rung-2 pilot: loop-only entrance-faded band LoRA, warm-start, inline eval; BW hard-only fallback

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

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)

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}

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