Spaces:
Sleeping
Sleeping
| """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 | |