src/peft/tuners/road/layer.py

# 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 warnings
from typing import Any, Optional, Union

import torch
import torch.nn as nn

from peft.tuners.tuners_utils import BaseTunerLayer, check_adapters_to_merge

from .config import RoadConfig, RoadVariant


class RoadLayer(BaseTunerLayer):
    """
    Road layer.

    Generally the idea of RoAD is to split the input vector into many 2D vectors and rotate each 2D vector with its own
    2D rotation matrix. For additional flexibility, each rotation matrix is multiplied by a trainable scale.

    when applied to vector R @ x each pair of elements of x is transformed like this: `y₀ = x₀ * α * cosθ - xₙ * α *
    sinθ` and `yₙ = x₀ * α * sinθ + xₙ * α * cosθ`

    The scales α and angles θ are learned for each pair of elements and, moreover, each of the 4 instances in the
    rotation matrix may actually be different (when using variant 2 or 4).

    Note that instead of using two consecutive elements x₀ x₁ we first split the whole vector into groups and pair
    elements from the first with the second half of the same group, which allows for more efficient inference
    implementation.

    The adapter needs to only store the angles θ and scales α, rather than the full matrix R and the inference
    implementation only needs to do elementwise vector multiplications.

    For merging the weights, we make use of the following formula: R @ (W @ x + b) = (R @ W) @ x + R @ b. The lhs part
    is how it is used in unmerged state (using efficient elementwise implementation instead of matrix multiplication)
    and the rhs part is how it is used in merged state where (R @ W) becomes the new weight matrix and R @ b becomes
    the new bias.

    """

    adapter_layer_names: tuple[str, ...] = ("road_theta", "road_alpha")
    other_param_names: tuple[str, ...] = ("variant", "group_size")

    def __init__(self, base_layer: nn.Module, ephemeral_gpu_offload: bool = False, **kwargs) -> None:
        self.base_layer = base_layer
        self.variant = {}
        self.group_size = {}
        self.road_theta = nn.ParameterDict({})
        self.road_alpha = nn.ParameterDict({})

        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
        else:
            raise ValueError(f"Unsupported layer type '{type(base_layer)}' encountered, cannot apply RoAd adapter.")
        self.in_features = in_features
        self.out_features = out_features

    @property
    def _available_adapters(self) -> set[str]:
        return {*self.road_theta}

    def update_layer(
        self,
        adapter_name,
        variant,
        group_size,
        init_weights,
        inference_mode: bool = False,
    ):
        self.variant[adapter_name] = variant
        self.group_size[adapter_name] = group_size

        if self.out_features % group_size != 0:
            raise ValueError(
                f"The out_features of the base layer must be divisible by group_size ({group_size}) when using RoadLayer."
            )

        # Actual trainable parameters
        if variant == "road_1":
            size = self.out_features // 2
        elif variant == "road_2":
            size = self.out_features
        elif variant == "road_4":
            size = self.out_features * 2
        else:
            raise ValueError(
                f"Unsupported variant {variant} for RoadLayer. Supported variants are road_1, road_2, and road_4."
            )
        self.road_theta[adapter_name] = nn.Parameter(torch.empty(size))
        self.road_alpha[adapter_name] = nn.Parameter(torch.empty(size))

        self.reset_parameters(adapter_name, init_weights)
        self._move_adapter_to_device_of_base_layer(adapter_name)

        self.set_adapter(self.active_adapters, inference_mode=inference_mode)

    def reset_parameters(self, adapter_name, init_weights):
        if init_weights is False:
            nn.init.normal_(self.road_theta[adapter_name].data, mean=0.0, std=0.5)
            nn.init.normal_(self.road_alpha[adapter_name].data, mean=1.0, std=0.5)
            return
        nn.init.zeros_(self.road_theta[adapter_name].data)
        nn.init.ones_(self.road_alpha[adapter_name].data)


class Linear(nn.Module, RoadLayer):
    # Road implemented in a dense layer
    def __init__(
        self,
        base_layer,
        adapter_name: str,
        variant: RoadVariant = "road_1",
        group_size: int = 64,
        init_weights: Union[bool, str] = True,
        **kwargs,
    ) -> None:
        super().__init__()
        RoadLayer.__init__(self, base_layer, **kwargs)

        self._active_adapter = adapter_name

        self.update_layer(
            adapter_name,
            variant,
            group_size,
            init_weights=init_weights,
        )

    def _check_forward_args(self, x, *args, **kwargs):
        """Check if the arguments are compatible with the configs and state of the model"""
        adapter_names = kwargs.get("adapter_names", None)
        if adapter_names is None:
            return

        if len(x) != len(adapter_names):
            msg = (
                "Length of `adapter_names` should be the same as the number of inputs, but got "
                f"{len(adapter_names)} and {len(x)} respectively."
            )
            raise ValueError(msg)

        if self.merged:
            # It is unclear what would be the right thing to do if users pass adapter_names and there are merged
            # adapters. Therefore, it is better to raise an error in this case.
            msg = "Cannot pass `adapter_names` when there are merged adapters, please call `unmerge_adapter` first."
            raise ValueError(msg)

    def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor:
        self._check_forward_args(x, *args, **kwargs)
        adapter_names = kwargs.pop("adapter_names", None)

        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)
        elif adapter_names is not None:
            result = self._mixed_batch_forward(x, *args, adapter_names=adapter_names, **kwargs)
        else:
            result = self.base_layer(x, *args, **kwargs)
            torch_result_dtype = result.dtype

            for active_adapter in self.active_adapters:
                if active_adapter not in self._available_adapters:
                    continue

                result = self._cast_input_dtype(result, self.road_theta[active_adapter].dtype)
                result = _apply_road(
                    self.variant[active_adapter],
                    self.group_size[active_adapter],
                    self.road_theta[active_adapter],
                    self.road_alpha[active_adapter],
                    result,
                )

            result = result.to(torch_result_dtype)

        return result

    def _mixed_batch_forward(
        self, x: torch.Tensor, *args: Any, adapter_names: list[str], **kwargs: Any
    ) -> torch.Tensor:
        # This is a special method that handles the case when users pass the argument `adapter_names`. This is an
        # extra argument that allows mixing different adapters in the same batch at inference time.
        result = self.base_layer(x, *args, **kwargs)

        unique_adapters = set(adapter_names)
        sub_batch_indices_list = []
        for adapter in unique_adapters:
            sub_batch_indices_list.append([index for index, item in enumerate(adapter_names) if item == adapter])

        for i, active_adapter in enumerate(unique_adapters):
            if active_adapter == "__base__":
                continue
            if active_adapter not in self._available_adapters:
                continue

            dtype = self.road_theta[active_adapter].data.dtype

            # getting the sub-batch, passing it to Road layers and updating the corresponding indices of the linear
            # layer output
            sub_batch = result[sub_batch_indices_list[i]].to(dtype)
            result[sub_batch_indices_list[i]] = _apply_road(
                self.variant[active_adapter],
                self.group_size[active_adapter],
                self.road_theta[active_adapter],
                self.road_alpha[active_adapter],
                sub_batch,
            )

        return result

    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`.
        """
        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._available_adapters:
                base_layer = self.get_base_layer()
                orig_dtype = base_layer.weight.dtype
                road_R = _get_delta_weight(
                    self.variant[active_adapter],
                    self.group_size[active_adapter],
                    self.road_theta[active_adapter].data,
                    self.road_alpha[active_adapter].data,
                )
                if safe_merge:
                    # Note that safe_merge will be slower than the normal merge
                    # because of the copy operation.
                    orig_weight = base_layer.weight.data.clone()
                    orig_weight = torch.matmul(road_R.to(orig_dtype), orig_weight)

                    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.contiguous().to(orig_dtype)

                    if base_layer.bias is not None:
                        orig_bias = base_layer.bias.clone()
                        orig_bias = torch.matmul(road_R.to(orig_dtype), orig_bias)

                        if not torch.isfinite(orig_bias).all():
                            raise ValueError(
                                f"NaNs detected in the merged bias. The adapter {active_adapter} seems to be broken"
                            )

                        base_layer.bias.data = orig_bias.contiguous().to(orig_dtype)
                else:
                    orig_weight = base_layer.weight.data
                    orig_weight = torch.matmul(road_R.to(orig_dtype), orig_weight)
                    base_layer.weight.data = orig_weight.contiguous().to(orig_dtype)

                    if base_layer.bias is not None:
                        orig_bias = base_layer.bias.data
                        orig_bias = torch.matmul(road_R.to(orig_dtype), orig_bias)
                        base_layer.bias.data = orig_bias.contiguous().to(orig_dtype)

                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:
            # Going in reverse order
            active_adapter = self.merged_adapters.pop()
            if active_adapter in self._available_adapters:
                weight = self.get_base_layer().weight
                orig_dtype = weight.dtype
                road_R = _get_delta_weight(
                    self.variant[active_adapter],
                    self.group_size[active_adapter],
                    self.road_theta[active_adapter].data,
                    self.road_alpha[active_adapter].data,
                )
                # Since our matrix are not necessarily orthogonal we need inverse instead of transpose.
                # In practice we expect this to basically always work since we start from block diagonal rotation matrix.
                inv_road_R = torch.linalg.inv(road_R.to(torch.float32)).to(orig_dtype)
                orig_weight = torch.matmul(inv_road_R, weight.data)
                weight.data = orig_weight.contiguous()

                if self.get_base_layer().bias is not None:
                    orig_bias = torch.matmul(inv_road_R, self.get_base_layer().bias.data)
                    self.get_base_layer().bias.data = orig_bias.contiguous()

    def __repr__(self) -> str:
        rep = super().__repr__()
        return "road." + rep


def _get_delta_weight(variant: RoadVariant, group_size: int, road_theta: torch.Tensor, road_alpha: torch.Tensor):
    first_col, second_col = _prepare_cols(variant, group_size, road_theta, road_alpha)

    # To help understand the logic below consider how rope embeddings work
    # here it is similar, but done in groups.
    # https://discuss.huggingface.co/t/is-llama-rotary-embedding-implementation-correct/44509/3

    # First column is simply put on the main diagonal
    output_tensor = torch.diag(first_col)
    # For second column we need to swap each half groups and add minus sign
    size = second_col.shape[0]
    swapped_second_col = second_col.reshape(-1, 2, group_size // 2)[:, [1, 0], :].flatten()
    rotated_diag_second_col = torch.diag(swapped_second_col).reshape(-1, 2, group_size // 2, size)[:, [1, 0], :, :]
    rotated_diag_second_col[:, 0, :, :] *= -1
    rotated_diag_second_col = rotated_diag_second_col.reshape(size, size)
    output_tensor += rotated_diag_second_col

    return output_tensor


def _prepare_cols(
    variant: RoadVariant, group_size: int, road_theta: torch.Tensor, road_alpha: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
    # In inference mode, this can be cached
    if variant == "road_1":
        # In each group there are only group_size // 2 parameters that are reused
        road_theta = road_theta.reshape(-1, group_size // 2).repeat_interleave(2, dim=0).flatten()
        road_alpha = road_alpha.reshape(-1, group_size // 2).repeat_interleave(2, dim=0).flatten()

        theta_cos = road_theta.cos()
        theta_sin = road_theta.sin()

        first_col = road_alpha * theta_cos
        second_col = road_alpha * theta_sin
    elif variant == "road_2":
        # Each group has exactly group_size parameters
        theta_cos = road_theta.cos()
        theta_sin = road_theta.sin()

        first_col = road_alpha * theta_cos
        second_col = road_alpha * theta_sin
    elif variant == "road_4":
        # Each group has 2*group_size parameters, first half used for first column, second half for second column
        road_theta = road_theta.reshape(-1, 2, group_size)
        theta_cos = road_theta[:, 0, :].cos().flatten()
        theta_sin = road_theta[:, 1, :].sin().flatten()
        road_alpha = road_alpha.reshape(-1, 2, group_size)
        alpha_1 = road_alpha[:, 0, :].flatten()
        alpha_2 = road_alpha[:, 1, :].flatten()

        first_col = alpha_1 * theta_cos
        second_col = alpha_2 * theta_sin
    else:
        raise ValueError(
            f"Unsupported variant {variant} for RoadLayer. Supported variants are road_1, road_2, and road_4."
        )

    return first_col, second_col


def _apply_road(
    variant: RoadVariant, group_size: int, road_theta: torch.Tensor, road_alpha: torch.Tensor, x: torch.Tensor
):
    first_col, second_col = _prepare_cols(variant, group_size, road_theta, road_alpha)

    # Split in half groups and join back
    # See equation 4 in the RoAD paper
    x_grouped = x.reshape(-1, 2, group_size // 2)
    x1 = x_grouped[:, 0, :]
    x2 = x_grouped[:, 1, :]
    rotate_half_x = torch.stack((-x2, x1), dim=1).reshape(x.shape)
    result = x * first_col + rotate_half_x * second_col
    return result


def dispatch_default(
    target: torch.nn.Module,
    adapter_name: str,
    road_config: RoadConfig,
    **kwargs,
) -> Optional[torch.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):
        new_module = Linear(target, adapter_name, **kwargs)

    return new_module