Skip to content

Low-Rank Adaptation (LoRA)

About

The d9d.peft.lora package implements Low-Rank Adaptation. It works by wrapping existing Linear layers (both standard nn.Linear and d9d's GroupedLinear) with a container that holds the original frozen layer (base) and two low-rank trainable matrices (lora_A and lora_B).

Because the original layer is moved to a submodule (base), the state keys change. The LoRA method automatically generates a ModelStateMapperRename to handle this transparently during checkpoint loading.

Usage Example

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
import torch
import re
from d9d.peft import inject_peft_and_freeze, merge_peft
from d9d.peft.lora import LoRA, LoRAConfig, LoRAParameters

# 1. Configuration
config = LoRAConfig(
    module_name_pattern=re.compile(r".*attention\.q_proj.*"),  # Target Attention Q projections
    params=LoRAParameters(
        r=16,
        alpha=32,
        dropout=0.1
    )
)

# 2. Instantiate Method
method = LoRA(config)

# 3. Inject
# This replaces nn.Linear with LoRALinear layers in-place.
# 'mapper' knows how to route 'q_proj.weight' -> 'q_proj.base.weight'
mapper = inject_peft_and_freeze(method, model)

# ... pass 'mapper' object to d9d's Trainer or manually load a model checkpoint ...

# ... train a model ...

# 4. Merge - for exporting a model
merge_peft(method, model)

d9d.peft.lora

Package for Low-Rank Adaptation (LoRA) implementation.

LoRA

Bases: PeftMethod[LoRAConfig]

Implements the Low-Rank Adaptation (LoRA) injection strategy.

It scans the module structure for nn.Linear or GroupedLinear layers matching the configured name pattern. Matched layers are replaced with LoRA wrappers.

It also generates ModelStateMapperRename objects. Since the original weight layer.weight is now at layer.base.weight inside the wrapper, the mapper ensures that loading a standard checkpoint still works by redirecting the key.

__init__(config)

Constructs a LoRA method.

Parameters:

Name Type Description Default
config LoRAConfig

LoRA configuration containing patterns and hyperparameters.

 required

LoRAConfig

Bases: BaseModel

Configuration for LoRA application.

Attributes:

Name Type Description
kind Literal['lora']

Discriminator field, always "lora".
module_name_pattern Pattern

Regular expression matching module names to wrap with LoRA.
params LoRAParameters

Hyperparameters for the LoRA layers.

LoRAGroupedLinear

Bases: Module

A LoRA wrapper around a GroupedLinear layer (commonly used in MoE or grouped query attention).

Attributes:

Name Type Description
lora_A

The A matrix (grouped linear).
lora_B

The B matrix (grouped linear).
base

The original base GroupedLinear layer.
dropout

Scaling dropout layer.

__init__(base_layer, params)

Constructs a LoRAGroupedLinear layer.

Parameters:

Name Type Description Default
base_layer GroupedLinear

The original GroupedLinear layer to wrap.

 required
params LoRAParameters

LoRA hyperparameters.

 required

forward(x, x_groups)

Computes forward pass for grouped inputs.

Parameters:

Name Type Description Default
x Tensor

Input tensor.

 required
x_groups Tensor

A tensor indicating group indices for each input.

 required

Returns:

Type Description
Tensor

Combined output of base and LoRA path.

merge_with_base_()

Collapse the LoRA weights into the base GroupedLinear layer.

Returns:

Type Description
GroupedLinear

The modified GroupedLinear layer.

reset_parameters()

Resets LoRA parameters. A is random, B is zeroed.

LoRALinear

Bases: Module

A LoRA wrapper around a standard PyTorch Linear layer.

Wraps a base linear layer and adds low-rank adaptation matrices A and B.

Attributes:

Name Type Description
lora_A

The A matrix (in_features -> r).
lora_B

The B matrix (r -> out_features).
base

The original base Linear layer.
dropout Dropout

Scaling dropout layer.

__init__(base_layer, params)

Constructs a LoRALinear layer.

Parameters:

Name Type Description Default
base_layer Linear

The original Linear layer to wrap.

 required
params LoRAParameters

LoRA hyperparameters (r, alpha, dropout).

 required

Raises:

Type Description
ValueError

If the base layer has a bias (currently unsupported).

forward(x)

Takes input tensor, computes base output and LoRA adaptation, and returns the sum.

Parameters:

Name Type Description Default
x Tensor

Input tensor.

 required

Returns:

Type Description
Tensor

The output of base(x) + scale * (B @ A @ dropout(x)).

merge_with_base_()

Collapse the LoRA weights into the base linear layer.

Returns:

Type Description
Linear

The modified base linear layer with updated weights.

reset_parameters()

Resets LoRA parameters. A is random, B is zeroed.

LoRAParameters

Bases: BaseModel

Hyperparameters for LoRA layers.

Attributes:

Name Type Description
r int

Rank of the low-rank adaptation matrices.
alpha int

Scaling factor for the learned weights.
dropout float

Dropout probability for the input to LoRA layers.