lora prep

conf/lora/lora-is-this-charlie-parker.yaml

@@ -0,0 +1,4 @@
+$include:
+  - conf/vampnet.yml
+
+fine_tune: True

scripts/exp/train.py

@@ -260,6 +260,7 @@ def train(
     suffix_amt: float = 0.0,
     prefix_dropout: float = 0.1,
     suffix_dropout: float = 0.1,
+    fine_tune: bool = False, 
     quiet: bool = False,
 ):
     assert codec_ckpt is not None, "codec_ckpt is required"
@@ -310,6 +311,11 @@ def train(
 
     criterion = CrossEntropyLoss()
 
+    if fine_tune:
+        import loralib as lora
+        lora.mark_only_lora_as_trainable(model)
+
+
     class Trainer(at.ml.BaseTrainer):
         _last_grad_norm = 0.0
 

vampnet/modules/transformer.py

@@ -5,6 +5,7 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from einops import rearrange
+import loralib as lora
 
 from .base import VampBase
 from .activations import get_activation
@@ -13,6 +14,8 @@ from .layers import FiLM
 from .layers import SequentialWithFiLM
 from .layers import WNConv1d
 
+LORA_R = 4
+
 
 class RMSNorm(nn.Module):
     def __init__(self, hidden_size: int, eps=1e-6):
@@ -86,9 +89,9 @@ class MultiHeadRelativeAttention(nn.Module):
         self.attention_max_distance = attention_max_distance
 
         # Create linear query, key, value projections
-        self.w_qs = nn.Linear(d_model, d_model, bias=False)
+        self.w_qs = lora.Linear(d_model, d_model, bias=False, r=LORA_R)
         self.w_ks = nn.Linear(d_model, d_model, bias=False)
-        self.w_vs = nn.Linear(d_model, d_model, bias=False)
+        self.w_vs = lora.Linear(d_model, d_model, bias=False, r=LORA_R)
 
         # Create linear final output projection
         self.fc = nn.Linear(d_model, d_model, bias=False)