From 0063e4356a8e8e7bb984a6d847afbb11d5fbaa4a Mon Sep 17 00:00:00 2001 From: hnluo <haoneng.lhn@alibaba-inc.com> Date: 星期五, 21 七月 2023 15:07:34 +0800 Subject: [PATCH] Merge pull request #768 from alibaba-damo-academy/dev_lhn --- funasr/modules/lora/layers.py | 323 +++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 files changed, 323 insertions(+), 0 deletions(-) diff --git a/funasr/modules/lora/layers.py b/funasr/modules/lora/layers.py new file mode 100644 index 0000000..76f046c --- /dev/null +++ b/funasr/modules/lora/layers.py @@ -0,0 +1,323 @@ +# ------------------------------------------------------------------------------------------ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License (MIT). See LICENSE in the repo root for license information. +# ------------------------------------------------------------------------------------------ +import torch +import torch.nn as nn +import torch.nn.functional as F + +import math +from typing import Optional, List + +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.: + self.lora_dropout = nn.Dropout(p=lora_dropout) + else: + self.lora_dropout = lambda x: x + # Mark the weight as unmerged + self.merged = False + self.merge_weights = merge_weights + + +class Embedding(nn.Embedding, LoRALayer): + # LoRA implemented in a dense layer + def __init__( + self, + num_embeddings: int, + embedding_dim: int, + r: int = 0, + lora_alpha: int = 1, + merge_weights: bool = True, + **kwargs + ): + nn.Embedding.__init__(self, num_embeddings, embedding_dim, **kwargs) + LoRALayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=0, + merge_weights=merge_weights) + # Actual trainable parameters + if r > 0: + self.lora_A = nn.Parameter(self.weight.new_zeros((r, num_embeddings))) + self.lora_B = nn.Parameter(self.weight.new_zeros((embedding_dim, r))) + self.scaling = self.lora_alpha / self.r + # Freezing the pre-trained weight matrix + self.weight.requires_grad = False + self.reset_parameters() + + def reset_parameters(self): + nn.Embedding.reset_parameters(self) + if hasattr(self, 'lora_A'): + # initialize A the same way as the default for nn.Linear and B to zero + nn.init.zeros_(self.lora_A) + nn.init.normal_(self.lora_B) + + def train(self, mode: bool = True): + nn.Embedding.train(self, mode) + if self.merge_weights and self.merged: + # Make sure that the weights are not merged + if self.r > 0: + self.weight.data -= (self.lora_B @ self.lora_A).T * self.scaling + self.merged = False + + def eval(self): + nn.Linear.eval(self) + if self.merge_weights and not self.merged: + # Merge the weights and mark it + if self.r > 0: + self.weight.data += (self.lora_B @ self.lora_A) * self.scaling + self.merged = True + + def forward(self, x: torch.Tensor): + if self.r > 0 and not self.merged: + result = nn.Embedding.forward(self, x) + if self.r > 0: + after_A = F.embedding( + x, self.lora_A.T, self.padding_idx, self.max_norm, + self.norm_type, self.scale_grad_by_freq, self.sparse + ) + result += (after_A @ self.lora_B.T) * self.scaling + return result + else: + return nn.Embedding.forward(self, x) + + +class Linear(nn.Linear, LoRALayer): + # LoRA implemented in a dense layer + def __init__( + self, + in_features: int, + out_features: int, + r: int = 0, + lora_alpha: int = 1, + lora_dropout: float = 0., + fan_in_fan_out: bool = False, # Set this to True if the layer to replace stores weight like (fan_in, fan_out) + merge_weights: bool = True, + **kwargs + ): + nn.Linear.__init__(self, in_features, out_features, **kwargs) + LoRALayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout, + merge_weights=merge_weights) + + self.fan_in_fan_out = fan_in_fan_out + # Actual trainable parameters + if r > 0: + self.lora_A = nn.Parameter(self.weight.new_zeros((r, in_features))) + self.lora_B = nn.Parameter(self.weight.new_zeros((out_features, r))) + self.scaling = self.lora_alpha / self.r + # Freezing the pre-trained weight matrix + self.weight.requires_grad = False + self.reset_parameters() + if fan_in_fan_out: + self.weight.data = self.weight.data.T + + def reset_parameters(self): + nn.Linear.reset_parameters(self) + if hasattr(self, 'lora_A'): + # initialize A the same way as the default for nn.Linear and B to zero + nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5)) + nn.init.zeros_(self.lora_B) + + def train(self, mode: bool = True): + def T(w): + return w.T if self.fan_in_fan_out else w + nn.Linear.train(self, mode) + if self.merge_weights and self.merged: + # Make sure that the weights are not merged + if self.r > 0: + self.weight.data -= T(self.lora_B @ self.lora_A) * self.scaling + self.merged = False + + def eval(self): + def T(w): + return w.T if self.fan_in_fan_out else w + nn.Linear.eval(self) + if self.merge_weights and not self.merged: + # Merge the weights and mark it + if self.r > 0: + self.weight.data += T(self.lora_B @ self.lora_A) * self.scaling + self.merged = True + + def forward(self, x: torch.Tensor): + def T(w): + return w.T if self.fan_in_fan_out else w + if self.r > 0 and not self.merged: + result = F.linear(x, T(self.weight), bias=self.bias) + if self.r > 0: + result += (self.lora_dropout(x) @ self.lora_A.T @ self.lora_B.T) * self.scaling + return result + else: + return F.linear(x, T(self.weight), bias=self.bias) + + +class MergedLinear(nn.Linear, LoRALayer): + # LoRA implemented in a dense layer + def __init__( + self, + in_features: int, + out_features: int, + r: int = 0, + lora_alpha: int = 1, + lora_dropout: float = 0., + enable_lora: List[bool] = [False], + fan_in_fan_out: bool = False, + merge_weights: bool = True, + **kwargs + ): + nn.Linear.__init__(self, in_features, out_features, **kwargs) + LoRALayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout, + merge_weights=merge_weights) + assert out_features % len(enable_lora) == 0, \ + 'The length of enable_lora must divide out_features' + self.enable_lora = enable_lora + self.fan_in_fan_out = fan_in_fan_out + # Actual trainable parameters + if r > 0 and any(enable_lora): + self.lora_A = nn.Parameter( + self.weight.new_zeros((r * sum(enable_lora), in_features))) + self.lora_B = nn.Parameter( + self.weight.new_zeros((out_features // len(enable_lora) * sum(enable_lora), r)) + ) # weights for Conv1D with groups=sum(enable_lora) + self.scaling = self.lora_alpha / self.r + # Freezing the pre-trained weight matrix + self.weight.requires_grad = False + # Compute the indices + self.lora_ind = self.weight.new_zeros( + (out_features, ), dtype=torch.bool + ).view(len(enable_lora), -1) + self.lora_ind[enable_lora, :] = True + self.lora_ind = self.lora_ind.view(-1) + self.reset_parameters() + if fan_in_fan_out: + self.weight.data = self.weight.data.T + + def reset_parameters(self): + nn.Linear.reset_parameters(self) + if hasattr(self, 'lora_A'): + # initialize A the same way as the default for nn.Linear and B to zero + nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5)) + nn.init.zeros_(self.lora_B) + + def zero_pad(self, x): + result = x.new_zeros((*x.shape[:-1], self.out_features)) + result = result.view(-1, self.out_features) + result[:, self.lora_ind] = x.reshape( + -1, self.out_features // len(self.enable_lora) * sum(self.enable_lora) + ) + return result.view((*x.shape[:-1], self.out_features)) + + def train(self, mode: bool = True): + def T(w): + return w.T if self.fan_in_fan_out else w + nn.Linear.train(self, mode) + if self.merge_weights and self.merged: + # Make sure that the weights are not merged + if self.r > 0 and any(self.enable_lora): + delta_w = F.conv1d( + self.lora_A.data.unsqueeze(0), + self.lora_B.data.unsqueeze(-1), + groups=sum(self.enable_lora) + ).squeeze(0) + self.weight.data -= self.zero_pad(T(delta_w * self.scaling)) + self.merged = False + + def eval(self): + def T(w): + return w.T if self.fan_in_fan_out else w + nn.Linear.eval(self) + if self.merge_weights and not self.merged: + # Merge the weights and mark it + if self.r > 0 and any(self.enable_lora): + delta_w = F.conv1d( + self.lora_A.data.unsqueeze(0), + self.lora_B.data.unsqueeze(-1), + groups=sum(self.enable_lora) + ).squeeze(0) + self.weight.data += self.zero_pad(T(delta_w * self.scaling)) + self.merged = True + + def forward(self, x: torch.Tensor): + def T(w): + return w.T if self.fan_in_fan_out else w + if self.merged: + return F.linear(x, T(self.weight), bias=self.bias) + else: + result = F.linear(x, T(self.weight), bias=self.bias) + if self.r > 0: + after_A = F.linear(self.lora_dropout(x), self.lora_A) + after_B = F.conv1d( + after_A.transpose(-2, -1), + self.lora_B.unsqueeze(-1), + groups=sum(self.enable_lora) + ).transpose(-2, -1) + result += self.zero_pad(after_B) * self.scaling + return result + + +class Conv2d(nn.Conv2d, LoRALayer): + # LoRA implemented in a dense layer + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + r: int = 0, + lora_alpha: int = 1, + lora_dropout: float = 0., + merge_weights: bool = True, + **kwargs + ): + nn.Conv2d.__init__(self, in_channels, out_channels, kernel_size, **kwargs) + LoRALayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout, + merge_weights=merge_weights) + assert type(kernel_size) is int + # Actual trainable parameters + if r > 0: + self.lora_A = nn.Parameter( + self.weight.new_zeros((r*kernel_size, in_channels*kernel_size)) + ) + self.lora_B = nn.Parameter( + self.weight.new_zeros((out_channels*kernel_size, r*kernel_size)) + ) + self.scaling = self.lora_alpha / self.r + # Freezing the pre-trained weight matrix + self.weight.requires_grad = False + self.reset_parameters() + + def reset_parameters(self): + nn.Conv2d.reset_parameters(self) + if hasattr(self, 'lora_A'): + # initialize A the same way as the default for nn.Linear and B to zero + nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5)) + nn.init.zeros_(self.lora_B) + + def train(self, mode: bool = True): + nn.Conv2d.train(self, mode) + if self.merge_weights and self.merged: + # Make sure that the weights are not merged + self.weight.data -= (self.lora_B @ self.lora_A).view(self.weight.shape) * self.scaling + self.merged = False + + def eval(self): + nn.Conv2d.eval(self) + if self.merge_weights and not self.merged: + # Merge the weights and mark it + self.weight.data += (self.lora_B @ self.lora_A).view(self.weight.shape) * self.scaling + self.merged = True + + def forward(self, x: torch.Tensor): + if self.r > 0 and not self.merged: + return F.conv2d( + x, + self.weight + (self.lora_B @ self.lora_A).view(self.weight.shape) * self.scaling, + self.bias, self.stride, self.padding, self.dilation, self.groups + ) + return nn.Conv2d.forward(self, x) + -- Gitblit v1.9.1