File size: 2,768 Bytes
b84549f |
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 |
import torch
from torch import nn
from abc import ABC, abstractmethod
from utils.dl.common.model import get_model_device, get_model_latency, get_model_size, set_module
from utils.common.log import logger
from .base import FMLoRA_Util, LoRA
class ToQKV_WrappedWithLoRA(nn.Module):
def __init__(self, fc: nn.Linear, ab_r: int):
super(ToQKV_WrappedWithLoRA, self).__init__()
self.fc = fc
self.ab = self.create_ab_as_linear(fc.weight.data, ab_r)
def create_ab_as_linear(self, fc_weight: torch.Tensor, ab_r: int):
res = nn.Sequential(
LoRA(fc_weight.size(1), fc_weight.size(0) // ab_r, bias=False),
LoRA(fc_weight.size(0) // ab_r, fc_weight.size(0), bias=False)
).to(fc_weight.device)
nn.init.kaiming_uniform_(res[0].weight, a=5 ** 0.5)
nn.init.zeros_(res[1].weight)
return res
def forward(self, x):
x1 = self.fc(x)
x2 = self.ab(x)
return x1 + x2
class FMLoRA_Vilt_Util(FMLoRA_Util):
@torch.no_grad()
def add_lora_ab_to_fm(self, fm: nn.Module, ab_r: int, samples: dict):
fm.eval()
# print(samples)
for k, v in samples.items():
if isinstance(v, torch.Tensor):
samples[k] = v.to(get_model_device(fm))
o1 = fm(**samples)
for name, module in fm.named_modules():
if name.endswith(('query', 'key', 'value')):
set_module(fm, name, ToQKV_WrappedWithLoRA(module, ab_r))
o2 = fm(**samples)
if isinstance(o1, tuple):
o1 = o1[-1]
o2 = o2[-1]
output_diff = ((o1.logits - o2.logits) ** 2).sum()
assert output_diff < 1e-5
return fm
@torch.no_grad()
def absorb_lora_and_recover_net_structure(self, fm: nn.Module, samples: dict):
fm.eval()
# print('absorb lora before')
for k, v in samples.items():
if isinstance(v, torch.Tensor):
samples[k] = v.to(get_model_device(fm))
o1 = fm(**samples)
for name, module in fm.named_modules():
if not isinstance(module, ToQKV_WrappedWithLoRA):
continue
fc = module.fc
ab = module.ab
fc.weight.add_(ab[1].weight @ ab[0].weight)
set_module(fm, name, fc)
# print('absorb lora after')
o2 = fm(**samples)
if isinstance(o1, tuple):
o1 = o1[-1]
o2 = o2[-1]
output_diff = ((o1.logits - o2.logits) ** 2).sum()
assert output_diff < 1e-6, output_diff
return fm
|