Spaces:
Runtime error
Runtime error
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import torch.nn.utils.parametrize as parametrize | |
| from typing import List | |
| class LoRALayer: | |
| def __init__(self, features_in: int, features_out: int, rank: int=1, alphas: int=1): | |
| super().__init__() | |
| self.lora_A = nn.Linear(features_in, rank, bias=False) | |
| self.lora_B = nn.Linear(rank, features_out, bias=False) | |
| nn.init.normal_(self.lora_A.weight, mean=0, std=1/rank) | |
| self.scale = alphas / rank | |
| class LoRALinear(nn.Module, LoRALayer): | |
| def __init__(self, base_layer: nn.Module, rank: int=1, alphas: int=1, dropout_p: float=0.0): | |
| features_out, features_in = base_layer.weight.shape | |
| super().__init__() | |
| LoRALayer.__init__(self, features_in=features_in, features_out=features_out, rank=rank, alphas=alphas) | |
| self.base_layer = nn.Linear(features_in, features_out, bias=False) | |
| self.base_layer.weight = base_layer.weight | |
| if dropout_p > 0.0: | |
| self.lora_dropout = nn.Dropout(p=dropout_p, inplace=False) | |
| else: | |
| self.lora_dropout = nn.Identity() | |
| self.enabled = False | |
| def forward(self, x: torch.Tensor): | |
| result = self.base_layer(x) | |
| if self.enabled: | |
| result = result + self.lora_B(self.lora_A(self.lora_dropout(x))) * self.scale | |
| return result | |
| def enable_lora(model: nn.Module, lora_modules=['q_proj', 'k_proj', 'v_proj', 'o_proj'], enabled=True): | |
| for name, module in model.named_modules(): | |
| if name.split('.')[-1] in lora_modules: | |
| module.enabled = enabled | |
| return model | |
| def replace_module(module: nn.Module, target_modules: List[str], torch_dtype: torch.dtype, **kwargs): | |
| for child_name, child_module in module.named_children(): | |
| if child_name in target_modules: | |
| new_module = LoRALinear(child_module, **kwargs).to(torch_dtype) | |
| setattr(module, child_name, new_module) | |
| else: | |
| replace_module(child_module, target_modules, torch_dtype, **kwargs) | |
| def get_lora_model(model: nn.Module, rank: float, alphas: float, lora_modules=['q_proj', 'k_proj', 'v_proj', 'o_proj'], dropout_p: float = 0.0, training: bool = False, torch_dtype: torch.dtype = torch.bfloat16): | |
| lora_config = {'rank': rank, | |
| 'alphas': alphas, | |
| 'dropout_p': dropout_p} | |
| replace_module(model, lora_modules, torch_dtype, **lora_config) | |
| for name, param in model.named_parameters(): | |
| if 'lora' not in name: | |
| param.requires_grad = False | |
| else: | |
| if training: | |
| param.requires_grad = True | |
| else: | |
| param.requires_grad = False | |
| return model | |