From ea4a7d324502176464cb26692254ca23d0007926 Mon Sep 17 00:00:00 2001
From: Eric Brown <ebrown5676@gmail.com>
Date: Fri, 22 Nov 2024 08:36:42 -0800
Subject: [PATCH] Add loraplus config

---
 src/fairseq2/models/llama/factory.py | 7 +++++++
 src/fairseq2/nn/lora.py              | 1 +
 2 files changed, 8 insertions(+)

diff --git a/src/fairseq2/models/llama/factory.py b/src/fairseq2/models/llama/factory.py
index 5b81ee538..b0a1ac14e 100644
--- a/src/fairseq2/models/llama/factory.py
+++ b/src/fairseq2/models/llama/factory.py
@@ -324,3 +324,10 @@ def get_llama_lora_config() -> LoRAConfig:
         dropout_p=0.05,
         keys=[".*decoder.layers.*.self_attn.*(q_proj|v_proj)$"],
     )
+
+
+def get_llama_loraplus_config(lr_ratio: float = 4.0) -> LoRAConfig:
+    config = get_llama_lora_config()
+    config.lr_ratio = lr_ratio
+
+    return config
diff --git a/src/fairseq2/nn/lora.py b/src/fairseq2/nn/lora.py
index acc98ef95..efe738887 100644
--- a/src/fairseq2/nn/lora.py
+++ b/src/fairseq2/nn/lora.py
@@ -30,6 +30,7 @@ class LoRAConfig:
     alpha: float
     dropout_p: float
     keys: list[str]
+    lr_ratio: float = 1.0
 
 
 class LoRALayer(ABC):