Class BitLora

bitnet/bit_lora.py:28–100

BitLora class represents a custom linear layer with LoRa (Low Rank) regularization. Args: rank (int): The rank of the LoRa regularization. Default is 4. lora_alpha (int): The scaling factor for LoRa regularization. Default is 1. *args: Variable length argument list.

class BitLora(BitLinear):
    """
    BitLora class represents a custom linear layer with LoRa (Low Rank) regularization.

    Args:
        rank (int): The rank of the LoRa regularization. Default is 4.
        lora_alpha (int): The scaling factor for LoRa regularization. Default is 1.
        *args: Variable length argument list.
        **kwargs: Arbitrary keyword arguments.

    Attributes:
        rank (int): The rank of the LoRa regularization.
        lora_alpha (int): The scaling factor for LoRa regularization.
        scaling (float): The scaling factor for LoRa regularization.
        merged (bool): Indicates whether the LoRa regularization has been merged with the weight matrix.
        lora_a (nn.Parameter): The learnable parameter matrix of shape (in_features, rank).
        lora_b (nn.Parameter): The learnable parameter matrix of shape (rank, out_features).

    Examples:

    """

    def __init__(self, rank: int = 4, lora_alpha: int = 1, *args, **kwargs):
        super(BitLora, self).__init__(*args, **kwargs)
        self.rank = rank
        self.lora_alpha = lora_alpha
        self.scaling = self.lora_alpha / self.rank
        self.merged = False

        self.lora_a = nn.Parameter(torch.zeros(self.in_features, rank))
        self.lora_b = nn.Parameter(torch.zeros(rank, self.out_features))

        # Rmsnorm
        self.rms_norm = SimpleRMSNorm(self.in_features)

    def forward(self, x: Tensor) -> Tensor:
        """
        Forward pass of the BitLora layer.

        Args:
            x (Tensor): The input tensor.

        Returns:
            Tensor: The output tensor.

        """
        w = self.weight

        # Normalize the input tensor
        x_norm = self.rms_norm(x)

        # Activation Quant
        x_quant = activation_quant(x_norm)

        if not self.merged and self.rank > 0:
            lora = self.lora_a @ self.lora_b
            w = w + lora * self.scaling