diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py index b885624..ed33ba7 100644 --- a/src/peft/peft_model.py +++ b/src/peft/peft_model.py @@ -195,7 +195,9 @@ class PeftModel(PushToHubMixin, torch.nn.Module): self.transformer_backbone_name = name if self.peft_config.num_transformer_submodules is None: - self.peft_config.num_transformer_submodules = 2 if self.peft_config.task_type == TaskType.SEQ_2_SEQ_LM else 1 + self.peft_config.num_transformer_submodules = ( + 2 if self.peft_config.task_type == TaskType.SEQ_2_SEQ_LM else 1 + ) for named_param, value in list(transformer_backbone.named_parameters()): if value.shape[0] == self.base_model.config.vocab_size: @@ -733,8 +735,9 @@ class PeftModelForSeq2SeqLM(PeftModel): decoder_inputs_embeds = torch.cat( (prompts[:, self.peft_config.num_virtual_tokens :], decoder_inputs_embeds), dim=1 ) - return self.base_model(inputs_embeds=inputs_embeds, decoder_inputs_embeds=decoder_inputs_embeds, **kwargs) - + return self.base_model( + inputs_embeds=inputs_embeds, decoder_inputs_embeds=decoder_inputs_embeds, **kwargs + ) def generate(self, **kwargs): if not isinstance(self.peft_config, PromptLearningConfig): diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index 6756b0c..ed10675 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -145,7 +145,7 @@ class LoraModel(torch.nn.Module): is_target_modules_in_base_model = True parent, target, target_name = self._get_submodules(key) bias = target.bias is not None - if loaded_in_8bit and isinstance(target, bnb.nn.Linear8bitLt) and self.peft_config.enable_lora is None: + if loaded_in_8bit and isinstance(target, bnb.nn.Linear8bitLt): kwargs.update( { "has_fp16_weights": target.state.has_fp16_weights, @@ -154,7 +154,11 @@ class LoraModel(torch.nn.Module): "index": target.index, } ) - new_module = Linear8bitLt(target.in_features, target.out_features, bias=bias, **kwargs) + if self.peft_config.enable_lora is None: + new_module = Linear8bitLt(target.in_features, target.out_features, bias=bias, **kwargs) + else: + kwargs.update({"enable_lora": self.peft_config.enable_lora}) + new_module = MergedLinear8bitLt(target.in_features, target.out_features, bias=bias, **kwargs) elif isinstance(target, torch.nn.Linear) and self.peft_config.enable_lora is None: new_module = Linear(target.in_features, target.out_features, bias=bias, **kwargs) elif self.peft_config.enable_lora is not None: @@ -509,3 +513,82 @@ if is_bnb_available(): output = self.lora_B(self.lora_A(self.lora_dropout(x))) * self.scaling result += output return result + + class MergedLinear8bitLt(bnb.nn.Linear8bitLt, 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.0, + enable_lora: List[bool] = [False], + **kwargs, + ): + bnb.nn.Linear8bitLt.__init__( + self, + in_features, + out_features, + bias=kwargs.get("bias", True), + has_fp16_weights=kwargs.get("has_fp16_weights", True), + memory_efficient_backward=kwargs.get("memory_efficient_backward", False), + threshold=kwargs.get("threshold", 0.0), + index=kwargs.get("index", None), + ) + LoraLayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout, merge_weights=False) + if out_features % len(enable_lora) != 0: + raise ValueError("The length of enable_lora must divide out_features") + self.enable_lora = enable_lora + # Actual trainable parameters + if r > 0 and any(enable_lora): + self.lora_A = nn.Linear(in_features, r * sum(enable_lora), bias=False) + self.lora_B = nn.Conv1d( + r * sum(enable_lora), + out_features // len(enable_lora) * sum(enable_lora), + kernel_size=1, + groups=2, + bias=False, + ) + 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() + + def 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.weight, a=math.sqrt(5)) + nn.init.zeros_(self.lora_B.weight) + + 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 forward(self, x: torch.Tensor): + result = super().forward(x) + if self.disable_adapters: + return result + elif self.r > 0: + if not torch.is_autocast_enabled(): + expected_dtype = result.dtype + if x.dtype != torch.float32: + x = x.float() + after_A = self.lora_A(self.lora_dropout(x)) + after_B = self.lora_B(after_A.transpose(-2, -1)).transpose(-2, -1) + output = self.zero_pad(after_B).to(expected_dtype) * self.scaling + result += output + else: + after_A = self.lora_A(self.lora_dropout(x)) + after_B = self.lora_B(after_A.transpose(-2, -1)).transpose(-2, -1) + output = self.zero_pad(after_B) * self.scaling + result += output + return result diff --git a/src/peft/utils/config.py b/src/peft/utils/config.py index cf1e49a..3e2cf5b 100644 --- a/src/peft/utils/config.py +++ b/src/peft/utils/config.py @@ -160,6 +160,8 @@ class PromptLearningConfig(PeftConfig): token_dim: int = field( default=None, metadata={"help": "The hidden embedding dimension of the base transformer model"} ) - num_transformer_submodules: Optional[int] = field(default=None, metadata={"help": "Number of transformer submodules"}) + num_transformer_submodules: Optional[int] = field( + default=None, metadata={"help": "Number of transformer submodules"} + ) num_attention_heads: Optional[int] = field(default=None, metadata={"help": "Number of attention heads"}) num_layers: Optional[int] = field(default=None, metadata={"help": "Number of transformer layers"})