From 53a8c9b609057056a31769d46c5a0cb916376ac0 Mon Sep 17 00:00:00 2001 From: Nils Date: Tue, 14 Jul 2026 10:29:57 +0200 Subject: [PATCH] adaptive-alpha merge (Lys-inspired, trained) + truncated-BPTT deep-k training Co-Authored-By: Claude Fable 5 --- scripts/eval_loop_code.py | 9 +++---- scripts/loop_common.py | 47 +++++++++++++++++++++++++++++++++++-- scripts/train_merge_code.py | 26 ++++++++++++++++---- 3 files changed, 72 insertions(+), 10 deletions(-) diff --git a/scripts/eval_loop_code.py b/scripts/eval_loop_code.py index b2dd04d..2cb6108 100644 --- a/scripts/eval_loop_code.py +++ b/scripts/eval_loop_code.py @@ -15,7 +15,7 @@ from pathlib import Path import torch -from loop_common import BandLooper, MergeAdapter +from loop_common import AdaptiveMergeAdapter, BandLooper, MergeAdapter from prep_mbpp import DIRECT_SUFFIX, extract_code, mbpp_prompt, run_tests sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) @@ -75,15 +75,16 @@ def main(): ap.add_argument("--pause", type=int, default=0, help="append p pause tokens to each prompt") ap.add_argument("--alpha", type=float, default=0.3) + ap.add_argument("--adaptive", action="store_true") args = ap.parse_args() ks = [int(x) for x in args.ks.split(",")] model, tok = load_model(dtype=torch.bfloat16) tok.padding_side = "left" looper = BandLooper(model) - adapter = MergeAdapter( - d=model.config.get_text_config().hidden_size, - alpha=args.alpha).cuda() + cls = AdaptiveMergeAdapter if args.adaptive else MergeAdapter + kw = ({"alpha0": args.alpha} if args.adaptive else {"alpha": args.alpha}) + adapter = cls(d=model.config.get_text_config().hidden_size, **kw).cuda() if args.adapter: adapter.load_state_dict(torch.load(args.adapter, map_location="cuda")) adapter.eval() diff --git a/scripts/loop_common.py b/scripts/loop_common.py index 0d9919c..294925a 100644 --- a/scripts/loop_common.py +++ b/scripts/loop_common.py @@ -53,6 +53,37 @@ class MergeAdapter(nn.Module): return out.to(dt) +class AdaptiveMergeAdapter(nn.Module): + """Merge with state-dependent anchor coefficient (Lys-inspired, trained). + + alpha(e, s) = sigmoid(w·[e;ŝ] + b), per position; w zero-init and + b = logit(0.3), so at init this is exactly the fixed alpha=0.3 merge.""" + + def __init__(self, d=1536, hidden=512, alpha0=0.3): + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(2 * d, hidden), nn.GELU(), nn.Linear(hidden, d) + ) + nn.init.zeros_(self.mlp[2].weight) + nn.init.zeros_(self.mlp[2].bias) + self.alpha_head = nn.Linear(2 * d, 1) + nn.init.zeros_(self.alpha_head.weight) + import math as _m + nn.init.constant_(self.alpha_head.bias, + _m.log(alpha0 / (1 - alpha0))) + + def forward(self, e, s): + dt = e.dtype + e32, s32 = e.float(), s.float() + s_hat = s32 * ( + e32.norm(dim=-1, keepdim=True) / (s32.norm(dim=-1, keepdim=True) + 1e-6) + ) + cat = torch.cat([e32, s_hat], dim=-1) + a = torch.sigmoid(self.alpha_head(cat)) + out = (1 - a) * e32 + a * s_hat + self.mlp(cat) + return out.to(dt) + + class BandLooper: """Capture layer-call kwargs once per forward, then re-run L14-30 manually.""" @@ -117,7 +148,9 @@ class BandLooper: def loop_logits(self, adapter, input_ids, k, attention_mask=None, use_checkpoint=False, return_states=False, last_only=False, - loop_mask=None, feedforward=False): + loop_mask=None, feedforward=False, bptt=None): + """bptt: backprop only through the last `bptt` iterations (McLeish- + style truncated BPTT); earlier iterations run under no_grad.""" """Teacher-forced logits after k merge->band loops. k=0 = plain forward. loop_mask (B, T) bool: positions where the merge applies; elsewhere the @@ -143,7 +176,17 @@ class BandLooper: with torch.no_grad(): s = self.band(e, calls) # s_0: no trainable params upstream states = [s] - for _ in range(k): + n_nograd = max(0, k - bptt) if bptt else 0 + for i in range(k): + if i < n_nograd: + with torch.no_grad(): + x = adapter(e, s) + if loop_mask is not None: + x = torch.where(loop_mask[..., None], x, e) + s = self.band(x, calls) + s = s.detach() + states.append(s) + continue x = adapter(e, s) if loop_mask is not None: x = torch.where(loop_mask[..., None], x, e) diff --git a/scripts/train_merge_code.py b/scripts/train_merge_code.py index 7eafcc6..b4de085 100644 --- a/scripts/train_merge_code.py +++ b/scripts/train_merge_code.py @@ -20,7 +20,7 @@ from pathlib import Path import torch import torch.nn.functional as F -from loop_common import BandLooper, MergeAdapter +from loop_common import AdaptiveMergeAdapter, BandLooper, MergeAdapter from prep_mbpp import DIRECT_SUFFIX, mbpp_prompt sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) @@ -34,6 +34,7 @@ LR = 1e-3 WARMUP = 20 MAX_TOK = 512 K_BUCKETS = [(1, ("easy",)), (2, ("easy", "hard")), (4, ("hard",))] +# --deepk 16 rescales to [(2, easy), (8, mixed), (16, hard)] ap = argparse.ArgumentParser() ap.add_argument("--seed", type=int, default=0) @@ -44,11 +45,19 @@ ap.add_argument("--feedforward", action="store_true", help="apply adapter once, no recurrence (pause-FF control)") ap.add_argument("--alpha", type=float, default=0.3, help="merge weight (2B-tuned default 0.3; try 0.1-0.15 at 12B)") +ap.add_argument("--adaptive", action="store_true", + help="state-dependent alpha (AdaptiveMergeAdapter)") +ap.add_argument("--deepk", type=int, default=0, + help="scale curriculum depths by deepk/4 (e.g. 16 -> 2/8/16)") +ap.add_argument("--bptt", type=int, default=0, + help="truncated BPTT: grads only through last N iterations") ARGS = ap.parse_args() SEED = ARGS.seed SUFFIX = ((f"_s{SEED}" if SEED else "") + (f"_p{ARGS.pause}" if ARGS.pause else "") - + (f"_a{ARGS.alpha}" if ARGS.alpha != 0.3 else "")) + + (f"_a{ARGS.alpha}" if ARGS.alpha != 0.3 else "") + + ("_ad" if ARGS.adaptive else "") + + (f"_dk{ARGS.deepk}" if ARGS.deepk else "")) PAUSE_ID = 6 # @@ -107,7 +116,15 @@ def main(): p.requires_grad_(False) looper = BandLooper(model) d = model.config.get_text_config().hidden_size - adapter = MergeAdapter(d=d, alpha=ARGS.alpha).cuda() + if ARGS.adaptive: + adapter = AdaptiveMergeAdapter(d=d, alpha0=ARGS.alpha).cuda() + else: + adapter = MergeAdapter(d=d, alpha=ARGS.alpha).cuda() + global K_BUCKETS + if ARGS.deepk: + f = ARGS.deepk / 4 + K_BUCKETS = [(max(1, int(k * f)), lbls) for k, lbls in K_BUCKETS] + print("K_BUCKETS ->", [(k, l) for k, l in K_BUCKETS], flush=True) opt = torch.optim.AdamW(adapter.parameters(), lr=LR, weight_decay=0.01) train = [it for it in data if it["split"] == "train" @@ -135,7 +152,8 @@ def main(): g["lr"] = lr_at(step) logits = looper.loop_logits(adapter, ids, k, attention_mask=msk, use_checkpoint=True, loop_mask=lmask, - feedforward=ARGS.feedforward) + feedforward=ARGS.feedforward, + bptt=ARGS.bptt or None) loss = F.cross_entropy(logits[:, :-1].flatten(0, 1).float(), lab[:, 1:].flatten(), ignore_index=-100) opt.zero_grad(set_to_none=True)