69 lines
2.3 KiB
Python
69 lines
2.3 KiB
Python
"""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}
|