File size: 3,985 Bytes
49dddf5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 | # Copyright 2026 Angshul Majumdar. Licensed under the Apache License, Version 2.0.
"""Discovery of adaptable weight matrices ('slots') and the forward hooks that apply dW."""
import re
from collections import defaultdict
import torch
import torch.nn as nn
try:
from transformers.pytorch_utils import Conv1D
except Exception: # pragma: no cover
class Conv1D(nn.Module):
pass
# GPT-2 family: Conv1D c_attn holds [q | k | v] as three equal column blocks
FUSED_EQUAL_SPLIT = {"c_attn": ("q", "k", "v")}
_LAYER_RE = re.compile(r"\.(\d+)\.")
class Slot:
"""One adapted matrix, always viewed as (d_in, d_out) so that the adapter adds x @ dW."""
def __init__(self, module_name, module, layer, role, c0, c1, d_in):
self.module_name, self.module = module_name, module
self.layer, self.role = layer, role
self.c0, self.c1, self.d_in = c0, c1, d_in
self.conv1d = isinstance(module, Conv1D)
@property
def shape(self):
return (self.d_in, self.c1 - self.c0)
def _view(self, w):
return w[:, self.c0:self.c1] if self.conv1d else w[self.c0:self.c1, :].T
def weight(self):
return self._view(self.module.weight)
def grad(self):
return self._view(self.module.weight.grad)
def __repr__(self):
return f"Slot(layer {self.layer}: {self.role} {list(self.shape)})"
def find_slots(model, target_modules="all-linear", exclude_modules=(), split_fused_qkv=True):
slots = []
for name, mod in model.named_modules():
if not isinstance(mod, (nn.Linear, Conv1D)):
continue
if any(e in name for e in exclude_modules):
continue
m = _LAYER_RE.search(name)
if m is None:
continue
role = name[m.end():]
w = mod.weight
d_in, d_out = (w.shape[0], w.shape[1]) if isinstance(mod, Conv1D) else (w.shape[1], w.shape[0])
parts = FUSED_EQUAL_SPLIT.get(role.split(".")[-1]) if (split_fused_qkv and isinstance(mod, Conv1D)) else None
if parts and d_out % len(parts) == 0:
wd = d_out // len(parts)
for j, p in enumerate(parts):
slots.append(Slot(name, mod, int(m.group(1)), f"{role}.{p}", j * wd, (j + 1) * wd, d_in))
else:
slots.append(Slot(name, mod, int(m.group(1)), role, 0, d_out, d_in))
if target_modules != "all-linear":
if isinstance(target_modules, str):
target_modules = [target_modules]
pats = [re.compile(r"(^|\.)" + re.escape(t) + r"$") for t in target_modules]
slots = [s for s in slots if any(p.search(s.role) for p in pats)]
if not slots:
raise ValueError(f"TensorPEFT: no target modules matched {target_modules!r}")
return slots
def attach_hooks(slots, delta_fn):
"""Forward hooks adding delta_fn(slot, x) to each slot's output columns (cast to the output dtype)."""
by_mod = defaultdict(list)
for s in slots:
by_mod[id(s.module)].append(s)
handles = []
for ss in by_mod.values():
ss = sorted(ss, key=lambda s: s.c0)
mod = ss[0].module
width = mod.weight.shape[1] if ss[0].conv1d else mod.weight.shape[0]
def hook(module, inputs, out, ss=ss, width=width):
x = inputs[0]
deltas = [delta_fn(s, x) for s in ss]
if all(d is None for d in deltas):
return out
pieces, pos = [], 0
for s, d in zip(ss, deltas):
if s.c0 > pos:
pieces.append(out.new_zeros(*out.shape[:-1], s.c0 - pos))
pieces.append(out.new_zeros(*out.shape[:-1], s.c1 - s.c0) if d is None else d.to(out.dtype))
pos = s.c1
if pos < width:
pieces.append(out.new_zeros(*out.shape[:-1], width - pos))
return out + (pieces[0] if len(pieces) == 1 else torch.cat(pieces, dim=-1))
handles.append(mod.register_forward_hook(hook))
return handles
|