xnn.common.finetune.lora.inject_lora#

xnn.common.finetune.lora.inject_lora(model, rank=4, alpha=1.0, init_std=0.001, *, exclude=(), freeze=True, trainable=())[source]#

Replace every eligible linear layer of model by a LoRA adapter, in place.

Eligible layers are torch.nn.Linear, e3nn.o3.Linear, the layers of e3nn.nn.FullyConnectedNet and every class registered with register_lora_target(). The adapted model computes exactly what model did until the adapters are trained.

Parameters:

model (torch.nn.Module) – The model (bare, multi-head or wrapped).
rank (int, optional) – Rank of every update, by default 4.
alpha (float, optional) – Scaling numerator, by default 1.0.
init_std (float, optional) – Standard deviation of the Gaussian factor at initialization.
exclude (sequence of str, optional) – fnmatch patterns of module names to leave alone (["readouts*"] keeps the readouts frozen and unadapted).
freeze (bool, optional) – Freeze every parameter that is not a LoRA factor (the usual LoRA setting), by default True.
trainable (sequence of str, optional) – Parameter-name patterns kept trainable besides the LoRA factors (["*atom_ref*"] lets the reference energies move).

Returns:

model.

Return type:

torch.nn.Module

Raises:

ValueError – If no layer was eligible.