Spaces:
Sleeping
Sleeping
File size: 3,691 Bytes
2320bb1 | 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 | """Matched-budget PEFT baselines for the demo (GPT-2), built on tensorpeft's slot discovery."""
import math
import torch
import torch.nn as nn
from tensorpeft.layers import attach_hooks, find_slots
class _Adapter(nn.Module):
lr = 1e-3
def trainable(self):
return [p for p in self.parameters() if p.requires_grad]
def n_params(self):
return sum(p.numel() for p in self.trainable())
def logits(self, ids):
return self.model[0](ids).logits
class LoRA(_Adapter):
lr = 1e-3
def __init__(self, model, rank=1, targets=("c_attn.q", "c_attn.v"), layers=None, seed=0):
super().__init__()
self.model = [model]
self.slots = [s for s in find_slots(model, list(targets)) if layers is None or s.layer in layers]
g = torch.Generator().manual_seed(seed)
dev = next(model.parameters()).device
self.A = nn.ParameterList([nn.Parameter(((torch.rand(s.shape[0], rank, generator=g) * 2 - 1)
/ math.sqrt(s.shape[0])).to(dev)) for s in self.slots])
self.B = nn.ParameterList([nn.Parameter(torch.zeros(rank, s.shape[1], device=dev)) for s in self.slots])
self.scale = 2.0
self._i = {id(s): i for i, s in enumerate(self.slots)}
self._h = attach_hooks(self.slots, lambda s, x: (x @ self.A[self._i[id(s)]]) @ self.B[self._i[id(s)]] * self.scale)
class IA3(_Adapter):
lr = 3e-3
def __init__(self, model, kv=("c_attn.k", "c_attn.v"), ff=True):
super().__init__()
self.model = [model]
self.out_slots = find_slots(model, list(kv))
self.ff_slots = find_slots(model, ["mlp.c_proj"]) if ff else []
dev = next(model.parameters()).device
self.l_out = nn.ParameterList([nn.Parameter(torch.ones(s.shape[1], device=dev)) for s in self.out_slots])
self.l_ff = nn.ParameterList([nn.Parameter(torch.ones(s.shape[0], device=dev)) for s in self.ff_slots])
io = {id(s): i for i, s in enumerate(self.out_slots)}
iff = {id(s): i for i, s in enumerate(self.ff_slots)}
def d_out(s, x):
y = x @ s.weight() + s.module.bias[s.c0:s.c1]
return y * (self.l_out[io[id(s)]] - 1)
def d_ff(s, x):
return (x * (self.l_ff[iff[id(s)]] - 1)) @ s.weight()
self._h = attach_hooks(self.out_slots, d_out) + attach_hooks(self.ff_slots, d_ff)
class Prompt(_Adapter):
lr = 3e-2
def __init__(self, model, n, seed=0):
super().__init__()
self.model = [model]
emb = model.get_input_embeddings().weight
g = torch.Generator().manual_seed(seed)
ids = torch.randint(0, emb.shape[0], (n,), generator=g).to(emb.device)
self.n = n
self.P = nn.Parameter(emb[ids].detach().clone())
def logits(self, ids):
m = self.model[0]
e = torch.cat([self.P.expand(ids.shape[0], -1, -1), m.get_input_embeddings()(ids)], dim=1)
return m(inputs_embeds=e).logits[:, self.n:]
class LNTune(_Adapter):
lr = 1e-3
def __init__(self, model, layers=None, final=True):
super().__init__()
self.model = [model]
self.params = []
for name, mod in model.named_modules():
if isinstance(mod, nn.LayerNorm):
parts = name.split(".")
layer = next((int(p) for p in parts if p.isdigit()), None)
if (layer is None and final) or (layer is not None and (layers is None or layer in layers)):
for p in mod.parameters():
p.requires_grad_(True)
self.params.append(p)
def trainable(self):
return self.params
|