Module neuralops.trainer.methods.lora

LoRA, DoRA, DisLoRA implementations.

Functions

def apply_dislora(model, config) ‑> None

Expand source code

@register_method("dislora")
def apply_dislora(model, config) -> None:
    """Apply dislora."""
    _apply_adapters(
        model,
        config.lora_r,
        config.lora_alpha,
        config.lora_dropout,
        "dislora",
        dislora_beta=config.dislora_beta,
        dislora_capacity_factor=config.dislora_capacity_factor,
    )

Apply dislora.

def apply_dora(model, config) ‑> None

Expand source code

@register_method("dora")
def apply_dora(model, config) -> None:
    """Apply dora."""
    _apply_adapters(model, config.lora_r, config.lora_alpha, config.lora_dropout, "dora")

Apply dora.

def apply_lora(model, config) ‑> None

Expand source code

@register_method("lora")
def apply_lora(model, config) -> None:
    """Apply lora."""
    _apply_adapters(model, config.lora_r, config.lora_alpha, config.lora_dropout, "lora")

Apply lora.

Classes

class DisLoRALinear (linear: torch.nn.modules.linear.Linear,
r: int,
alpha: float,
dropout: float,
beta: float = 0.5,
capacity_factor: float = 1.0)

Expand source code

class DisLoRALinear(nn.Module):
    def __init__(
        self,
        linear: nn.Linear,
        r: int,
        alpha: float,
        dropout: float,
        beta: float = 0.5,
        capacity_factor: float = 1.0,
    ) -> None:
        super().__init__()
        self.linear = linear
        self.linear.weight.requires_grad = False
        if linear.bias is not None:
            linear.bias.requires_grad = False
        self.r = r
        self.beta = beta
        self.capacity_factor = capacity_factor
        if r > 0:
            self.lora_A = nn.Linear(linear.in_features, r, bias=False)
            self.lora_B = nn.Linear(r, linear.out_features, bias=False)
            kaiming_uniform_(self.lora_A.weight, a=math.sqrt(5))
            zeros_(self.lora_B.weight)
            self.scaling = alpha / r
            self.dropout = nn.Dropout(dropout)
            self.register_buffer("_eye", torch.eye(r))
        else:
            self.lora_A = self.lora_B = None
            self.scaling = 1.0
            self.dropout = nn.Identity()

    def forward(self, x) -> Any:
        """Forward pass — compute module output from inputs."""
        base = self.linear(x)
        if self.r > 0:
            lora = self.lora_B(self.lora_A(self.dropout(x))) * self.scaling
            return base + lora
        return base

    def ortho_loss(self) -> Any:
        """Ortho loss."""
        if self.lora_A is None:
            return 0.0
        A = self.lora_A.weight
        B = self.lora_B.weight
        reg = (A @ A.T - self._eye).pow(2).sum() + (B.T @ B - self._eye).pow(2).sum()
        return self.beta * reg

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• torch.nn.modules.module.Module

Methods

def forward(self, x) ‑> Any

Expand source code

def forward(self, x) -> Any:
    """Forward pass — compute module output from inputs."""
    base = self.linear(x)
    if self.r > 0:
        lora = self.lora_B(self.lora_A(self.dropout(x))) * self.scaling
        return base + lora
    return base

Forward pass — compute module output from inputs.

def ortho_loss(self) ‑> Any

Expand source code

def ortho_loss(self) -> Any:
    """Ortho loss."""
    if self.lora_A is None:
        return 0.0
    A = self.lora_A.weight
    B = self.lora_B.weight
    reg = (A @ A.T - self._eye).pow(2).sum() + (B.T @ B - self._eye).pow(2).sum()
    return self.beta * reg

Ortho loss.

class DoRALinear (linear: torch.nn.modules.linear.Linear, r: int, alpha: float, dropout: float)

Expand source code

class DoRALinear(nn.Module):
    def __init__(self, linear: nn.Linear, r: int, alpha: float, dropout: float) -> None:
        """Execute __init__ operation."""
        super().__init__()
        self.linear = linear
        self.linear.weight.requires_grad = False
        if linear.bias is not None:
            linear.bias.requires_grad = False
        self.r = r
        if r > 0:
            self.m = nn.Parameter(torch.ones(linear.out_features))
            self.lora_A = nn.Linear(linear.in_features, r, bias=False)
            self.lora_B = nn.Linear(r, linear.out_features, bias=False)
            kaiming_uniform_(self.lora_A.weight, a=math.sqrt(5))
            zeros_(self.lora_B.weight)
            self.scaling = alpha / r
            self.dropout = nn.Dropout(dropout)
        else:
            self.m = nn.Parameter(torch.ones(linear.out_features))
            self.lora_A = self.lora_B = None
            self.scaling = 1.0
            self.dropout = nn.Identity()

    def forward(self, x) -> Any:
        """Forward pass — compute module output from inputs."""
        base_out = self.linear(x)
        if self.r > 0:
            lora = self.lora_B(self.lora_A(self.dropout(x))) * self.scaling
            weight_norm = self.linear.weight.norm(p=2, dim=1, keepdim=True)
            direction = self.linear.weight / (weight_norm + 1e-8)
            delta = self.m.unsqueeze(1) * direction + lora
            return F.linear(x, delta, self.linear.bias)
        return base_out

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

Execute init operation.

Ancestors

• torch.nn.modules.module.Module

Methods

def forward(self, x) ‑> Any

Expand source code

def forward(self, x) -> Any:
    """Forward pass — compute module output from inputs."""
    base_out = self.linear(x)
    if self.r > 0:
        lora = self.lora_B(self.lora_A(self.dropout(x))) * self.scaling
        weight_norm = self.linear.weight.norm(p=2, dim=1, keepdim=True)
        direction = self.linear.weight / (weight_norm + 1e-8)
        delta = self.m.unsqueeze(1) * direction + lora
        return F.linear(x, delta, self.linear.bias)
    return base_out

Forward pass — compute module output from inputs.