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>
2026-07-14 11:00:42 +02:00
parent 507c170744
commit 8fe8e34b47
3 changed files with 235 additions and 1 deletions
@@ -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}