loraify_linears_in_selection

penzai.toolshed.lora.loraify_linears_in_selection(selection: pz.Selection[Any], rank: int) → Any

Replaces Linear layers inside a selected part of a model with LoRA blocks.

This function should usually be called after freezing the existing weights in the model using something like

pz.nn.at_instances_of(pz.nn.Parameter).apply(
    lambda param: pz.nn.FrozenParameter(param.value, param.name)
)

This function returns a copy of the model with new LoRA parameters added, but does not modify any existing parameters.

Parameters:

• selection – A selection of a model that identifies the parts for which LoRA adaptation should be applied. Any Linear layers contained within the selected part will be replaced.

• rank – The rank of the LoRA blocks to insert.

Returns:

A copy of the original full model (e.g. of selection.deselect()), but where each of the Linear layers inside the selected part are replaced with new LowRankAdapter instances. Each such instance will contain UninitializedParameter instances that need to be initialized before use; their names will be based on the location of the Linear blocks in the original tree.