From 37c5185faea2553de206b2fb2b0cdacad4a52718 Mon Sep 17 00:00:00 2001
From: 51616 <rujikorn.ch@gmail.com>
Date: Mon, 13 Jan 2025 19:07:23 +0000
Subject: [PATCH] allow modules_to_save (for trainable layernorm)

---
 hyperlora/configs.py       | 4 ++++
 hyperlora/model_loading.py | 5 +++--
 2 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/hyperlora/configs.py b/hyperlora/configs.py
index 8526615..da42059 100644
--- a/hyperlora/configs.py
+++ b/hyperlora/configs.py
@@ -240,6 +240,10 @@ class LoRAArguments:
         default=None,
         metadata={"help": ("LoRA target modules.")},
     )
+    modules_to_save: Optional[list[str]] = field(
+        default=None,
+        metadata={"help": ("Modules to save.")},
+    )
 
 
 @dataclass
diff --git a/hyperlora/model_loading.py b/hyperlora/model_loading.py
index 0f7da38..48740a2 100644
--- a/hyperlora/model_loading.py
+++ b/hyperlora/model_loading.py
@@ -142,8 +142,9 @@ def get_model(
     if peft_config is not None:
         model = PeftModel(model, peft_config)
     model.train(train)
-    for param in model.parameters():
-        param.requires_grad = requires_grad
+    for name, param in model.named_parameters():
+        if "modules_to_save" not in name:
+            param.requires_grad = requires_grad
     return model