# Copyright 2025-present the HuggingFace Inc. team. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import math import warnings from typing import Optional import torch import torch.nn as nn from transformers.pytorch_utils import Conv1D from peft.tuners.tuners_utils import BaseTunerLayer from peft.utils.other import transpose class GraloraLayer(BaseTunerLayer): # List all names of layers that may contain adapter weight adapter_layer_names = ("gralora_A", "gralora_B", "gralora_A_general", "gralora_B_general") other_param_names = ("r", "hybrid_r", "alpha", "scaling", "gralora_dropout") def __init__(self, base_layer: nn.Module, **kwargs): self.base_layer = base_layer self.r = {} self.alpha = {} self.gralora_k = {} self.hybrid_r = {} self.scaling = {} self.gralora_dropout = nn.ModuleDict({}) self.gralora_A = nn.ParameterDict({}) self.gralora_B = nn.ParameterDict({}) self.gralora_A_general = nn.ModuleDict({}) self.gralora_B_general = nn.ModuleDict({}) # Mark the weight as unmerged self._disable_adapters = False self.merged_adapters = [] base_layer = self.get_base_layer() if isinstance(base_layer, nn.Linear): in_features, out_features = base_layer.in_features, base_layer.out_features elif isinstance(base_layer, Conv1D): in_features, out_features = ( base_layer.weight.ds_shape if hasattr(base_layer.weight, "ds_shape") else base_layer.weight.shape ) else: raise NotImplementedError(f"Unsupported layer type {type(base_layer)}") self.in_features = in_features self.out_features = out_features self.kwargs = kwargs def update_layer( self, adapter_name, module_name, r, alpha, gralora_dropout, gralora_k: int = 2, hybrid_r: int = 0, init_weights: bool = True, ): if r <= 0: raise ValueError(f"`r` should be a positive integer value but the value passed is {r}") elif hybrid_r < 0: raise ValueError(f"`hybrid_r` should be a non-negative integer value but the value passed is {hybrid_r}") self.r[adapter_name] = r self.alpha[adapter_name] = alpha self.gralora_k[adapter_name] = gralora_k self.hybrid_r[adapter_name] = hybrid_r if gralora_dropout > 0.0: gralora_dropout_layer = nn.Dropout(p=gralora_dropout) else: gralora_dropout_layer = nn.Identity() self.gralora_dropout.update(nn.ModuleDict({adapter_name: gralora_dropout_layer})) # Actual trainable parameters if self.in_features % gralora_k != 0: raise ValueError( f"in_features should be divisible by gralora_k, but got {self.in_features} and {gralora_k}" ) if self.out_features % gralora_k != 0: raise ValueError( f"out_features should be divisible by gralora_k, but got {self.out_features} and {gralora_k}" ) subblock_in_features = self.in_features // gralora_k subblock_out_features = self.out_features // gralora_k # gralora_r is the rank allocated to GraLoRA method; hybrid_r is the rank allocated to vanilla LoRA gralora_r = r gralora_A = [] gralora_B = [] for _ in range(gralora_k): new_A = nn.Parameter(torch.empty(gralora_r, subblock_in_features)) new_B = nn.Parameter(torch.empty(subblock_out_features, gralora_r)) if init_weights: # Initialize to identity: A is random, B is zero nn.init.kaiming_uniform_(new_A, a=math.sqrt(5)) nn.init.zeros_(new_B) else: # Initialize to random: both A and B are random (for testing) nn.init.kaiming_uniform_(new_A, a=math.sqrt(5)) nn.init.kaiming_uniform_(new_B, a=math.sqrt(5)) gralora_A.append(new_A) gralora_B.append(new_B) # stack A and B and transpose to get the final shape gralora_A = torch.stack(tuple(gralora_A), dim=0) # [N, gralora_r, in_features//N] gralora_A = gralora_A.transpose(1, 2).contiguous() # [N, in_features//N, gralora_r] gralora_B = torch.stack(tuple(gralora_B), dim=0) # [N, out_features//N, gralora_r] gralora_B = gralora_B.transpose(1, 2).contiguous() # [N, gralora_r, out_features//N] if hybrid_r > 0: general_gralora_A = nn.Linear(self.in_features, hybrid_r, bias=False) general_gralora_B = nn.Linear(hybrid_r, self.out_features, bias=False) if init_weights: # Initialize to identity: A is random, B is zero nn.init.kaiming_uniform_(general_gralora_A.weight, a=math.sqrt(5)) nn.init.zeros_(general_gralora_B.weight) else: # Initialize to random: both A and B are random (for testing) nn.init.kaiming_uniform_(general_gralora_A.weight, a=math.sqrt(5)) nn.init.kaiming_uniform_(general_gralora_B.weight, a=math.sqrt(5)) else: general_gralora_A = nn.Identity() general_gralora_B = nn.Identity() self.gralora_A[adapter_name] = gralora_A self.gralora_B[adapter_name] = gralora_B self.gralora_A_general[adapter_name] = general_gralora_A self.gralora_B_general[adapter_name] = general_gralora_B self.module_name = module_name self.scaling[adapter_name] = alpha / (gralora_r + hybrid_r) self._move_adapter_to_device_of_base_layer(adapter_name) self.set_adapter(self.active_adapters) class Linear(nn.Linear, GraloraLayer): # Gralora implemented in a dense layer def __init__( self, base_layer, adapter_name: str, module_name, r: int = 0, alpha: int = 1, gralora_dropout: float = 0.0, gralora_k: int = 2, hybrid_r: int = 0, fan_in_fan_out: bool = False, # Set this to True if the layer to replace stores weight like (fan_in, fan_out) init_weights: bool = True, **kwargs, ) -> None: # this gets the init from nn.Linear's super perspective, i.e. nn.Module.__init__, which should always be called super(nn.Linear, self).__init__() GraloraLayer.__init__(self, base_layer, **kwargs) self.fan_in_fan_out = fan_in_fan_out self._active_adapter = adapter_name self.update_layer(adapter_name, module_name, r, alpha, gralora_dropout, gralora_k, hybrid_r, init_weights) def merge(self, safe_merge: bool = False, adapter_names: Optional[list[str]] = None) -> None: """ Merge the active adapter weights into the base weights Args: safe_merge (`bool`, *optional*): If True, the merge operation will be performed in a copy of the original weights and check for NaNs before merging the weights. This is useful if you want to check if the merge operation will produce NaNs. Defaults to `False`. adapter_names (`list[str]`, *optional*): The list of adapter names that should be merged. If None, all active adapters will be merged. Defaults to `None`. """ from peft.tuners.tuners_utils import check_adapters_to_merge adapter_names = check_adapters_to_merge(self, adapter_names) if not adapter_names: # no adapter to merge return for active_adapter in adapter_names: if active_adapter in self.gralora_A.keys(): base_layer = self.get_base_layer() if safe_merge: # Note that safe_merge will be slower than the normal merge # because of the copy operation. orig_weights = base_layer.weight.data.clone() delta_weight = self.get_delta_weight(active_adapter) orig_weights += delta_weight if not torch.isfinite(orig_weights).all(): raise ValueError( f"NaNs detected in the merged weights. The adapter {active_adapter} seems to be broken" ) base_layer.weight.data = orig_weights else: delta_weight = self.get_delta_weight(active_adapter) base_layer.weight.data += delta_weight self.merged_adapters.append(active_adapter) def unmerge(self) -> None: """ This method unmerges all merged adapter layers from the base weights. """ if not self.merged: warnings.warn("Already unmerged. Nothing to do.") return while len(self.merged_adapters) > 0: active_adapter = self.merged_adapters.pop() if active_adapter in self.gralora_A.keys(): delta_weight = self.get_delta_weight(active_adapter) self.get_base_layer().weight.data -= delta_weight def get_delta_weight(self, adapter) -> torch.Tensor: """ Compute the delta weight for GraLoRA adapter. GraLoRA applies block-wise low-rank adaptation with information exchange. This method computes the equivalent weight matrix that would be added to the base weight during merge. Args: adapter (str): The name of the adapter Returns: torch.Tensor: The delta weight matrix with shape [out_features, in_features] """ gralora_A = self.gralora_A[adapter] # [N, in_features//N, rank] gralora_B = self.gralora_B[adapter] # [N, rank, out_features//N] gralora_A_general = self.gralora_A_general[adapter] gralora_B_general = self.gralora_B_general[adapter] device = gralora_A.device dtype = gralora_A.dtype gralora_k = self.gralora_k[adapter] hybrid_r = self.hybrid_r[adapter] r = self.r[adapter] # Handle CPU fp16/bf16 casting cast_to_fp32 = device.type == "cpu" and (dtype == torch.float16 or dtype == torch.bfloat16) if cast_to_fp32: gralora_A = gralora_A.float() gralora_B = gralora_B.float() # Get dimensions in_features = self.in_features out_features = self.out_features gralora_rank = r subblock_gralora_rank = gralora_rank // gralora_k # scatter gralora_A to get the scattered weight matrix l_indices = torch.arange(in_features, device=device) n_indices = l_indices // (in_features // gralora_k) i_indices = l_indices % (in_features // gralora_k) gralora_A_scattered = torch.zeros( in_features, gralora_k, gralora_rank, device=device, dtype=torch.float32 if cast_to_fp32 else dtype ) gralora_A_scattered.scatter_( 1, n_indices.unsqueeze(1).unsqueeze(2).expand(-1, 1, gralora_rank), gralora_A[n_indices, i_indices, :].unsqueeze(1), ) # compute the delta weight delta_weight = ( torch.einsum( "ikr, kro -> iko", gralora_A_scattered.view(in_features, gralora_k, gralora_k, subblock_gralora_rank) .permute(0, 2, 1, 3) .reshape(in_features, gralora_k, gralora_rank), gralora_B, ) .reshape(in_features, out_features) .T ) # Add hybrid LoRA component if present if hybrid_r > 0: weight_A_general = gralora_A_general.weight # [hybrid_r, in_features] weight_B_general = gralora_B_general.weight # [out_features, hybrid_r] if cast_to_fp32: weight_A_general = weight_A_general.float() weight_B_general = weight_B_general.float() # Compute delta for hybrid part: [out_features, hybrid_r] @ [hybrid_r, in_features] delta_weight += weight_B_general @ weight_A_general # Apply scaling and transpose if needed delta_weight = transpose(delta_weight, self.fan_in_fan_out) * self.scaling[adapter] # Cast back if needed if cast_to_fp32: delta_weight = delta_weight.to(dtype=dtype) return delta_weight def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor: previous_dtype = x.dtype if self.disable_adapters: if self.merged: self.unmerge() result = self.base_layer(x, *args, **kwargs) elif self.merged: result = self.base_layer(x, *args, **kwargs) else: result = self.base_layer(x, *args, **kwargs) torch_result_dtype = result.dtype # Handle 2D input: [batch, features] -> [batch, 1, features] # This is common for MLPs and other non-sequence models x_is_2d = x.ndim == 2 if x_is_2d: x = x.unsqueeze(1) # [B, F] -> [B, 1, F] for active_adapter in self.active_adapters: if active_adapter not in self.gralora_A.keys(): continue gralora_A = self.gralora_A[active_adapter] gralora_B = self.gralora_B[active_adapter] gralora_A_general = self.gralora_A_general[active_adapter] gralora_B_general = self.gralora_B_general[active_adapter] r = self.r[active_adapter] gralora_rank = r gralora_k = self.gralora_k[active_adapter] hybrid_r = self.hybrid_r[active_adapter] dropout = self.gralora_dropout[active_adapter] scaling = self.scaling[active_adapter] gralora_dtype = gralora_A.dtype B, L, in_features = x.shape N = gralora_k subblock_gralora_rank = gralora_rank // N output = torch.einsum( "bljr, jro -> bljo", torch.einsum( "blni, nir -> blnr", dropout(x.to(gralora_dtype)).view(B, L, N, in_features // N), gralora_A, ) .view(B, L, N, N, subblock_gralora_rank) .permute(0, 1, 3, 2, 4) .reshape(B, L, N, N * subblock_gralora_rank), gralora_B, ).reshape(B, L, -1) # Squeeze back to 2D if input was 2D if x_is_2d: output = output.squeeze(1) # [B, 1, F] -> [B, F] result += scaling * output.to(torch_result_dtype) if hybrid_r > 0: hybrid_output = gralora_B_general(gralora_A_general(dropout(x.to(gralora_dtype)))) if x_is_2d: hybrid_output = hybrid_output.squeeze(1) result += scaling * hybrid_output.to(torch_result_dtype) result = result.to(previous_dtype) return result def supports_lora_conversion(self, adapter_name: str = "default") -> bool: return True def __repr__(self) -> str: rep = super().__repr__() return "gralora." + rep