Skip to content

LoRA

About

The d9d.peft.lora package implements Low-Rank Adaptation. It wraps linear layers, both nn.Linear and d9d's GroupedLinear. The wrapper holds the original frozen layer (base) and two trainable low-rank layers (lora_A and lora_B). The low-rank layers of a GroupedLinear use the same grouped GEMM backend as the base layer. LoRA does not support nn.Linear layers with a bias.

State Mapping

Because the original layer moves to the base submodule, its state keys change. LoRA returns a ModelStateMapperRename for each wrapped layer. Chain these mappers after your checkpoint mapper, so standard checkpoints still load (see the example below).

Usage

 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
30
31
32
33
34
35
36
37
38
39
import re

from d9d.model_state.mapper.adapters import identity_mapper_from_module
from d9d.model_state.mapper.compose import ModelStateMapperSequential
from d9d.peft import inject_peft_and_freeze, merge_peft
from d9d.peft.lora import LoRA, LoRAConfig, LoRAParameters

model = ...  # Your model, before its weights are loaded.

# Maps the checkpoint keys to the original model structure.
checkpoint_mapper = identity_mapper_from_module(model)

# 1. Configure LoRA for the attention query projections.
config = LoRAConfig(
    module_name_pattern=re.compile(r".*self_attn\.q_proj"),
    params=LoRAParameters(
        r=16,
        alpha=32,
        dropout=0.1
    )
)

# 2. Create the method.
method = LoRA(config)

# 3. Inject.
# Replaces the matching nn.Linear layers with LoRALinear layers in place.
# The PEFT mapper only routes "q_proj.weight" to "q_proj.base.weight".
peft_mapper = inject_peft_and_freeze(method, model)

# Keys that the PEFT mapper does not change pass through unchanged.
mapper = ModelStateMapperSequential([checkpoint_mapper, peft_mapper])

# ... return the mapper from your ModelProvider, or load a checkpoint with it ...

# ... train the model ...

# 4. Merge the adapters into the base weights before exporting the model.
merge_peft(method, model)

API Reference

d9d.peft.lora

Low-rank adaptation (LoRA).

LoRA

Bases: PeftMethod[LoRAConfig]

PEFT method for Low-Rank Adaptation (LoRA).

inject replaces every nn.Linear or GroupedLinear layer whose name matches the configured pattern with a LoRA wrapper. The original layer.weight moves to layer.base.weight. For each wrapped layer, inject returns a ModelStateMapperRename, so standard checkpoints still load.

__init__(config)

Constructs the LoRA object.

Parameters:

Name Type Description Default
config LoRAConfig

LoRA configuration with the module name pattern 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 that must fully match the names of modules to wrap with LoRA.
params LoRAParameters

Hyperparameters for the LoRA layers.

LoRAGroupedLinear

Bases: Module

Wrapper that adds low-rank adaptation matrices A and B to a GroupedLinear layer, as used by MoE experts.

Attributes:

Name Type Description
lora_A

The A matrix (in_features -> r for each group).
lora_B

The B matrix (r -> out_features for each group).
base

The original GroupedLinear layer.
dropout

Dropout applied to the input of the LoRA path.
__init__(base_layer, params)

Constructs the LoRAGroupedLinear object.

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 the base output plus the scaled LoRA update for grouped inputs.

Parameters:

Name Type Description Default
x Tensor

Input tokens of all groups, with the tokens of each group stored together. Shape: (num_tokens, in_features).

 required
x_groups ExpertTokenGroups

How the tokens are split between groups.

 required

Returns:

Type Description
Tensor

The sum of the base and LoRA outputs. Shape: (num_tokens, out_features).
merge_with_base_()

Merges the LoRA weights into the base GroupedLinear layer in place.

Returns:

Type Description
GroupedLinear

The base GroupedLinear layer with updated weights.
reset_parameters()

Resets the LoRA parameters.

lora_A gets a random initialization and lora_B is zeroed, so the LoRA update starts at zero.

LoRALinear

Bases: Module

Wrapper that adds low-rank adaptation matrices A and B to an nn.Linear layer.

Attributes:

Name Type Description
lora_A

The A matrix (in_features -> r).
lora_B

The B matrix (r -> out_features).
base

The original nn.Linear layer.
dropout Dropout

Dropout applied to the input of the LoRA path.
__init__(base_layer, params)

Constructs the LoRALinear object.

Parameters:

Name Type Description Default
base_layer Linear

The original nn.Linear layer to wrap.

 required
params LoRAParameters

LoRA hyperparameters.

 required

Raises:

Type Description
ValueError

If base_layer has a bias.
forward(x)

Computes the base output plus the scaled LoRA update.

Parameters:

Name Type Description Default
x Tensor

Input tensor.

 required

Returns:

Type Description
Tensor

base(x) + alpha / r * lora_B(lora_A(dropout(x))).
merge_with_base_()

Merges the LoRA weights into the base linear layer in place.

Returns:

Type Description
Linear

The base linear layer with updated weights.
reset_parameters()

Resets the LoRA parameters.

lora_A gets a random initialization and lora_B is zeroed, so the LoRA update starts at zero.

LoRAParameters

Bases: BaseModel

Configuration for LoRA layers.

Attributes:

Name Type Description
r int

Rank of the low-rank adaptation matrices.
alpha int

Scaling factor for the learned weights. The LoRA output is scaled by alpha / r.
dropout float

Dropout probability for the input to LoRA layers.