File size: 4,455 Bytes
96ba80e | 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 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 | import math
from torch import Tensor
import torch.nn.functional as F
import torch
import torch.nn as nn
from quantize.int_linear import QuantLinear
def identity(x):
return x
class LoRALayer:
def __init__(
self,
r: int,
lora_alpha: int,
lora_dropout: float,
merge_weights: bool,
):
self.r = r
self.lora_alpha = lora_alpha
# Optional dropout
if lora_dropout > 0.0:
self.lora_dropout = nn.Dropout(p=lora_dropout)
else:
self.lora_dropout = identity
# Mark the weight as unmerged
self.merged = False
self.merge_weights = merge_weights
class LoRAQuantLinear(QuantLinear, LoRALayer):
"""
Quantized Module that can perform quantized convolution or normal convolution.
To activate quantization, please use set_quant_state function.
"""
def __init__(
self,
org_module: nn.Linear,
weight_quant_params: dict = {},
act_quant_params: dict = {},
disable_input_quant=False,
r=0,
lora_alpha=1,
lora_dropout=0.0,
merge_weights=True,
lora_attr={
"lora_iter_num" : 1,
"lora_quant" : False,
"lora_r":4,
"lora_only":False,
},
):
super().__init__(
org_module, weight_quant_params, act_quant_params, disable_input_quant
)
LoRALayer.__init__(
self,
r=r,
lora_alpha=lora_alpha,
lora_dropout=lora_dropout,
merge_weights=merge_weights,
)
self.lora_iter_num = lora_attr["lora_iter_num"]
self.lora_quant = lora_attr["lora_quant"]
self.lora_only = lora_attr["lora_only"]
if "lora_r" in lora_attr:
r = lora_attr["lora_r"]
self.r = lora_attr["lora_r"]
# Freezing the pre-trained weight matrix
self.weight.requires_grad = False
if self.r >0 : #sign_lora is in weight_quantizer
out_features, in_features = self.weight.shape
self.lora_A = nn.ParameterList([nn.Parameter(self.weight.new_zeros((r, in_features))) for _ in range(self.lora_iter_num)])
self.lora_B = nn.ParameterList([nn.Parameter(self.weight.new_zeros((out_features, r))) for _ in range(self.lora_iter_num)])
self.scaling = self.lora_alpha / r
self.reset_lora_parameters()
self.rms_norm = None
def update_quant_parms(self,weight_quant_params):
for k,v in weight_quant_params.items():
self.weight_quantizer.__setattr__(k,v)
def reset_lora_parameters(self):
if hasattr(self, "lora_A"):
# initialize A the same way as the default for nn.Linear and B to zero
for i in range(self.lora_iter_num):
nn.init.kaiming_uniform_(self.lora_A[i], a=math.sqrt(5))
nn.init.zeros_(self.lora_B[i])
def forward(self, input: torch.Tensor):
if self.use_temporary_parameter:
weight = self.temp_weight
bias = self.temp_bias
else:
weight = self.weight
bias = self.bias
if weight.device != input.device:
weight = weight.to(input.device)
if bias is not None:
bias = bias.to(input.device)
if self.merged:
weight = weight
else:
if self.use_temporary_parameter or self.use_weight_quant:
if self.r > 0:
weight = self.weight_quantizer(weight + self.lora_B[0] @ self.lora_A[0] * self.scaling, self.quant_rate)
else:
weight = self.weight_quantizer(weight, self.quant_rate)
else:
weight = weight
if self.use_act_quant and not self.disable_input_quant:
input = self.act_quantizer(input,self.quant_rate)
out = self.fwd_func(input, weight, bias, **self.fwd_kwargs)
return out
def extra_repr(self):
s = super().extra_repr()
s += ", use_temporary_parameter={}".format(self.use_temporary_parameter)
s += ", use_act_quant={}".format(self.use_act_quant)
s += ", use_weight_quant={}".format(self.use_weight_quant)
s += ", disable_input_quant={}".format(self.disable_input_quant)
# s += ", lora_quant"
return s
|