# Copyright 2026-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. from __future__ import annotations import warnings from typing import Any, Optional import torch from torch import nn, svd_lowrank from peft.tuners._buffer_dict import BufferDict from peft.tuners.tuners_utils import BaseTunerLayer, _get_in_out_features, check_adapters_to_merge from peft.utils.integrations import gather_params_ctx from peft.utils.other import transpose from .config import PsoftConfig class OrthLayer(nn.Module): """ r*r orthogonal transformation R used in PSOFT between A and B. Forward: output = input @ R.T """ def __init__( self, size: int, orth: bool = True, mag_b: bool = True, mag_a: bool = True, use_cayley_neumann: bool = False, num_cayley_neumann_terms: int = 5, cayley_neumann_eps: Optional[float] = None, ): super().__init__() self.size = size self.orth = orth self.mag_b = mag_b self.mag_a = mag_a self.use_cayley_neumann = use_cayley_neumann self.num_cayley_neumann_terms = num_cayley_neumann_terms self.cayley_neumann_eps = cayley_neumann_eps if orth: self.weight = nn.Parameter(torch.empty((size * (size - 1)) // 2)) rows, cols = torch.triu_indices(size, size, 1) self.register_buffer("rows", rows, persistent=False) self.register_buffer("cols", cols, persistent=False) else: self.weight = nn.Parameter(torch.empty(size, size)) self.vector_b = nn.Parameter(torch.empty(size)) if mag_b else None self.vector_a = nn.Parameter(torch.empty(size)) if mag_a else None def reset_parameters(self, init_weights: bool = True) -> None: params = [self.weight] if self.vector_b is not None: params.append(self.vector_b) if self.vector_a is not None: params.append(self.vector_a) if any(p.is_meta for p in params): return with torch.no_grad(): if init_weights: if self.orth: self.weight.zero_() else: nn.init.eye_(self.weight) if self.vector_b is not None: self.vector_b.fill_(1.0) if self.vector_a is not None: self.vector_a.fill_(1.0) else: if self.orth: nn.init.normal_(self.weight, mean=0.0, std=0.1) else: nn.init.eye_(self.weight) self.weight.add_(torch.randn_like(self.weight) * 0.1) if self.vector_b is not None: self.vector_b.fill_(1.0) if self.vector_a is not None: self.vector_a.fill_(1.0) def forward(self, input: torch.Tensor) -> torch.Tensor: R = self.get_matrix() if input.device.type == "cpu" and input.dtype in (torch.float16, torch.bfloat16): compute_dtype = torch.float32 else: compute_dtype = input.dtype return (input.to(compute_dtype) @ R.to(compute_dtype).t()).to(input.dtype) # Adapted from the Cayley/Neumann-based orthogonal parametrization used in OFT v2 # (PEFT implementation: https://github.com/huggingface/peft/blob/main/src/peft/tuners/oft/layer.py) #L104 def _skew_symmetric(self) -> torch.Tensor: Q = torch.zeros((self.size, self.size), device=self.weight.device, dtype=self.weight.dtype) Q = Q.index_put((self.rows, self.cols), self.weight) return Q - Q.transpose(0, 1) # Adapted from the Cayley/Neumann-based orthogonal parametrization used in OFT v2 # (PEFT implementation: https://github.com/huggingface/peft/blob/main/src/peft/tuners/oft/layer.py) #L160 def _project_Q(self, Q: torch.Tensor, eps: float = 0.9) -> torch.Tensor: norm = torch.linalg.norm(Q, ord="fro") if torch.isfinite(norm) and norm > eps: Q = Q * (eps / (norm + 1e-12)) return Q # R = (I+Q)(I-Q)^(-1) def get_matrix(self) -> torch.Tensor: cast_to_fp32 = False orig_dtype = None if not self.orth: R = self.weight else: Q = self._skew_symmetric() orig_dtype = Q.dtype id_mat = torch.eye(self.size, device=Q.device, dtype=Q.dtype) if self.use_cayley_neumann: if self.cayley_neumann_eps is not None: Q = self._project_Q(Q, eps=self.cayley_neumann_eps) t = int(self.num_cayley_neumann_terms) R = id_mat.clone() if t > 1: R.add_(Q, alpha=2.0) if t > 2: Q_squared = Q @ Q R.add_(Q_squared, alpha=2.0) Q_power = Q_squared for _ in range(3, t - 1): Q_power = Q_power @ Q R.add_(Q_power, alpha=2.0) Q_power = Q_power @ Q R.add_(Q_power) else: cast_to_fp32 = orig_dtype in (torch.float16, torch.bfloat16) if cast_to_fp32: Q = Q.float() # solver requires float32 R = torch.linalg.solve(id_mat - Q, id_mat + Q, left=False) # Apply scaling vectors to R if self.vector_b is not None: R = self.vector_b[:, None] * R if self.vector_a is not None: R = R * self.vector_a[None, :] if cast_to_fp32: R = R.to(orig_dtype) return R def __repr__(self) -> str: return ( f"psoft.{self.__class__.__name__}(" f"size={self.size}, orth={self.orth}, " f"use_cayley_neumann={self.use_cayley_neumann}, " f"num_cayley_neumann_terms={int(self.num_cayley_neumann_terms)}, " f"cayley_neumann_eps={self.cayley_neumann_eps}, " f"mag_a={self.mag_a}, mag_b={self.mag_b}" f")" ) class PsoftLayer(BaseTunerLayer): adapter_layer_names: tuple[str, ...] = ("psoft_R",) other_param_names: tuple[str, ...] = ( "r", "psoft_alpha", "scaling", "psoft_dropout", "psoft_svd", "psoft_svd_lowrank_niter", "ab_svd_init", ) def __init__(self, base_layer: nn.Module, **kwargs) -> None: super().__init__() self.base_layer = base_layer # per-adapter hyperparams self.r: dict[str, int] = {} self.psoft_alpha: dict[str, float] = {} self.scaling: dict[str, float] = {} self.psoft_dropout = nn.ModuleDict({}) self.psoft_svd: dict[str, str] = {} self.psoft_svd_lowrank_niter: dict[str, int] = {} self.ab_svd_init: dict[str, Optional[str]] = {} # per-adapter trainable module self.psoft_R = nn.ModuleDict({}) # per-adapter cache state self._psoft_A_cache = BufferDict(persistent=False) self._psoft_B_cache = BufferDict(persistent=False) self.merged_adapters: list[str] = [] self._disable_adapters = False self.kwargs = kwargs self.fan_in_fan_out = False base_layer = self.get_base_layer() in_features, out_features = _get_in_out_features(base_layer) self.in_features = in_features self.out_features = out_features def _get_psoft_ab_cache_buffers(self, adapter_name: str): return self._psoft_A_cache[adapter_name], self._psoft_B_cache[adapter_name] def _set_psoft_ab_cache_buffers(self, adapter_name: str, A: torch.Tensor, B: torch.Tensor) -> None: self._psoft_A_cache[adapter_name] = A self._psoft_B_cache[adapter_name] = B def update_layer(self, adapter_name: str, config: PsoftConfig, **kwargs: Any) -> None: ab_svd_init = config.ab_svd_init init_weights = config.init_weights r = int(config.r) self.fan_in_fan_out = config.fan_in_fan_out self.r[adapter_name] = r self.psoft_alpha[adapter_name] = config.psoft_alpha self.scaling[adapter_name] = config.psoft_alpha / r self.psoft_dropout[adapter_name] = ( nn.Dropout(p=config.psoft_dropout) if config.psoft_dropout > 0.0 else nn.Identity() ) self.ab_svd_init[adapter_name] = config.ab_svd_init self.psoft_svd[adapter_name] = config.psoft_svd self.psoft_svd_lowrank_niter[adapter_name] = config.psoft_svd_lowrank_niter self.psoft_R[adapter_name] = OrthLayer( size=r, orth=config.psoft_orth, mag_b=config.psoft_mag_b, mag_a=config.psoft_mag_a, use_cayley_neumann=config.use_cayley_neumann, num_cayley_neumann_terms=config.num_cayley_neumann_terms, cayley_neumann_eps=config.cayley_neumann_eps, ) self._move_adapter_to_device_of_base_layer(adapter_name) self.psoft_R[adapter_name].reset_parameters(init_weights=init_weights) self.psoft_R[adapter_name].requires_grad_(True) with gather_params_ctx(self.get_base_layer().weight): self._build_psoft_ab_cache_buffers(adapter_name, ab_svd_init) self.set_adapter([adapter_name]) # Adapted from the asymmetric SVD used in PiSSA # (PEFT implementation: https://github.com/huggingface/peft/blob/main/src/peft/tuners/lora/layer.py) #L316 def _build_psoft_ab_cache_buffers(self, adapter_name: str, init_type: str) -> None: with torch.no_grad(): base = self.get_base_layer() weight = base.weight dtype = weight.dtype if dtype not in (torch.float32, torch.float16, torch.bfloat16): raise TypeError("PSOFT init requires float32/float16/bfloat16. Re-quantize after init if needed.") # W: (out, in) fp32 W = transpose(weight.to(torch.float32), self.fan_in_fan_out) r = self.r[adapter_name] Vr, Sr, Uhr = self._compute_svd_factors( W, r, svd_mode=self.psoft_svd[adapter_name], niter=self.psoft_svd_lowrank_niter[adapter_name], ) Sr_scaled = Sr / self.scaling[adapter_name] if init_type == "psoft_init": A = Uhr # (r, in) B = Vr @ torch.diag(Sr_scaled) # (out, r) elif init_type == "pissa_init": s_sqrt = torch.sqrt(Sr_scaled) A = torch.diag(s_sqrt) @ Uhr # (r, in) B = Vr @ torch.diag(s_sqrt) # (out, r) else: raise ValueError(f"Unknown ab_svd_init: {init_type}") A = A.contiguous().detach() B = B.contiguous().detach() self._set_psoft_ab_cache_buffers(adapter_name, A, B) def _compute_svd_factors(self, weight: torch.Tensor, r: int, *, svd_mode: str, niter: int): # weight: (out, in) fp32 if svd_mode == "full": U, S, Vh = torch.linalg.svd(weight.data, full_matrices=False) Vr = U[:, :r] # (out, r) Sr = S[:r] # (r,) Uhr = Vh[:r, :] # (r, in) elif svd_mode == "lowrank": U, S, V = svd_lowrank(weight.data, q=r, niter=niter) # V: (in, r) Vr = U[:, :r] Sr = S[:r] Uhr = V[:, :r].t() # (r, in) else: raise ValueError(f"Unknown svd_mode: {svd_mode}") return Vr, Sr, Uhr class Linear(nn.Module, PsoftLayer): def __init__( self, base_layer: nn.Module, adapter_name: str, config: PsoftConfig, **kwargs: Any, ) -> None: super().__init__() PsoftLayer.__init__(self, base_layer, **kwargs) self.fan_in_fan_out = config.fan_in_fan_out self._active_adapter = adapter_name self.update_layer(adapter_name, config=config, **kwargs) def _get_R_matrix(self, adapter_name: str) -> torch.Tensor: return self.psoft_R[adapter_name].get_matrix() def get_delta_weight(self, adapter_name: str) -> torch.Tensor: """ ΔW = scaling * B (R - id_mat) A Returns in base weight layout (respecting fan_in_fan_out). """ A, B = self._get_psoft_ab_cache_buffers(adapter_name) base_w = self.get_base_layer().weight device = base_w.device out_dtype = base_w.dtype R = self._get_R_matrix(adapter_name) r = self.r[adapter_name] compute_dtype = ( torch.float32 if (device.type == "cpu" and out_dtype in (torch.float16, torch.bfloat16)) else out_dtype ) A_c = A.to(device=device, dtype=compute_dtype) B_c = B.to(device=device, dtype=compute_dtype) R_c = R.to(device=device, dtype=compute_dtype) id_mat = torch.eye(r, device=device, dtype=compute_dtype) delta = B_c @ (R_c - id_mat) @ A_c # (out, in) delta = transpose(delta, self.fan_in_fan_out) delta = delta * self.scaling[adapter_name] return delta.to(dtype=out_dtype) def merge(self, safe_merge: bool = False, adapter_names: Optional[list[str]] = None) -> None: adapter_names = check_adapters_to_merge(self, adapter_names) if not adapter_names: return base_layer = self.get_base_layer() for active_adapter in adapter_names: if active_adapter not in self.psoft_R: continue if safe_merge: orig_weight = base_layer.weight.data.clone() orig_dtype = orig_weight.dtype delta_weight = self.get_delta_weight(active_adapter) orig_weight += delta_weight.to(orig_dtype) if not torch.isfinite(orig_weight).all(): raise ValueError( f"NaNs detected in the merged weights. The adapter {active_adapter} seems to be broken" ) base_layer.weight.data = orig_weight else: delta_weight = self.get_delta_weight(active_adapter) base_layer.weight.data += delta_weight.to(base_layer.weight.dtype) self.merged_adapters.append(active_adapter) def supports_lora_conversion(self, adapter_name: str = "default") -> bool: return True def unmerge(self) -> None: if not self.merged: warnings.warn("Already unmerged. Nothing to do.", UserWarning) return weight = self.get_base_layer().weight while len(self.merged_adapters) > 0: active_adapter = self.merged_adapters.pop() if active_adapter not in self.psoft_R: continue orig_dtype = weight.dtype delta_weight = self.get_delta_weight(active_adapter) weight.data -= delta_weight.to(orig_dtype) def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor: 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 psoft_keys = self.psoft_R.keys() for active_adapter in self.active_adapters: if active_adapter not in psoft_keys: continue A, B = self._get_psoft_ab_cache_buffers(active_adapter) dropout = self.psoft_dropout[active_adapter] scaling = self.scaling[active_adapter] R_layer = self.psoft_R[active_adapter] x_cast = self._cast_input_dtype(x, A.dtype) x_d = dropout(x_cast) A_c = A.to(device=x_d.device, dtype=x_d.dtype) B_c = B.to(device=x_d.device, dtype=x_d.dtype) xa = x_d @ A_c.t() xr = R_layer(xa) delta_y = (xr - xa) @ B_c.t() result = result + (delta_y * scaling) result = result.to(torch_result_dtype) return result def __repr__(self) -> str: return "psoft." + super().__repr__() def dispatch_default( target: nn.Module, adapter_name: str, config: PsoftConfig, **kwargs, ) -> Optional[nn.Module]: new_module = None if isinstance(target, BaseTunerLayer): target_base_layer = target.get_base_layer() else: target_base_layer = target if isinstance(target_base_layer, torch.nn.Linear): if config.fan_in_fan_out: warnings.warn( "fan_in_fan_out=True is not compatible with `torch.nn.Linear`. Setting fan_in_fan_out=False." ) config.fan_in_fan_out = False new_module = Linear(target, adapter_name, config=config, **kwargs) return new_module