File size: 2,343 Bytes
8e04e6f | 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 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 | import loralib as lora
from torch import nn
def replace_lora_layers(module, r, lora_alpha, lora_dropout):
"""
Recursively replace all nn.Linear and nn.Embedding layers in the module with the lora.Linear and lora.Embedding layers.
Args:
module (nn.Module): The module containing nn.Linear layers to be replaced.
r (int): The rank for the pair of low-rank adaptation matrices.
lora_alpha (float): Used for calculating lora scaling.
lora_dropout (float): Dropout rate for the input before LoRA layers.
"""
for name, child in module.named_children():
if "pair_update" in name:
continue
# Check if the child module is an instance of nn.Linear
if isinstance(child, nn.Linear):
# Replace the nn.Linear layer with a new one
in_features = child.in_features
out_features = child.out_features
bias = child.bias is not None
# Turn off merge weights, as turning it on behaves strangely in the training mode of lightning trainer
new_layer = lora.Linear(
in_features,
out_features,
r,
lora_alpha,
lora_dropout,
merge_weights=False,
bias=bias,
)
setattr(module, name, new_layer)
elif isinstance(child, nn.Embedding):
# Replace nn.Embedding layer
num_embeddings = child.num_embeddings
embedding_dim = child.embedding_dim
padding_idx = child.padding_idx
max_norm = child.max_norm
norm_type = child.norm_type
scale_grad_by_freq = child.scale_grad_by_freq
sparse = child.sparse
new_layer = lora.Embedding(
num_embeddings,
embedding_dim,
r,
lora_alpha,
merge_weights=False,
padding_idx=padding_idx,
max_norm=max_norm,
norm_type=norm_type,
scale_grad_by_freq=scale_grad_by_freq,
sparse=sparse,
)
setattr(module, name, new_layer)
else:
# Recursively replace layers in submodules
replace_lora_layers(child, r, lora_alpha, lora_dropout)
|