xnn.common.finetune.lora.LoRAAdapter

class xnn.common.finetune.lora.LoRAAdapter(base, rank=4, alpha=1.0, *, weight='weight', in_axis=1, init_std=0.001)

Bases: Module

Low-rank update of one dense weight matrix of any module.

The base module is evaluated unchanged with its weight replaced by W + (alpha / rank) * delta, delta = lora_in @ lora_out arranged in the weight’s layout; lora_in (input side, Gaussian) has shape (d_in, rank) and lora_out (output side, zero at the start) (rank, d_out). The base weight is frozen.

Parameters:

• base (torch.nn.Module) – The module to adapt.

• rank (int, optional) – Rank of the update, by default 4.

• alpha (float, optional) – Scaling numerator, by default 1.0.

• weight (str, optional) – Name of the 2-D weight parameter, by default "weight".

• in_axis (int, optional) – Which axis of the weight is the input dimension: 1 for torch.nn.Linear ((d_out, d_in)), 0 for the (d_in, d_out) layout of e3nn’s fully connected layers. By default 1.

• init_std (float, optional) – Standard deviation of lora_in at initialization, by default 1e-3.

Variables:

• base (torch.nn.Module) – The adapted module (frozen weight).

• lora_out (lora_in,) – The two factors.

• scaling (float) – alpha / rank.

delta()

The unscaled update, in the layout of the base weight.

Return type:

Tensor

adapted_weight()

W + scaling * delta.

Return type:

Tensor

forward(*args, **kwargs)

Evaluate the base module with the adapted weight.

merge()

Fold the update into the base weight and return the base module.

Return type:

Module

extra_repr()

Set the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

Return type:

str