Repository: huggingface/peft

Issue:
'LoraModel.disable_adapter_layers' not causing the model to use the original modules
I am doing a unit test for another project that does this :
- Create a Lora model using the argument modules_to_save to train one additional layer (e.g., `base_model.model.transformer.h.3.mlp.c_proj` for gpt2)
- Make some logit predictions with the initial model
- Do some training
- Disable the Lora adapter with the method disable_adapter_layers
- Do some new logit predictions, and verify that the results are the same as before the training

This test is failing, what I see is that while the Lora layers are indeed bypassed due to the call to LoraModel.disable_adapter_layers, the layer given in modules_to_save is used with its post-training parameters, not its original parameters. The forward pass uses the layer `base_model.model.transformer.h.3.mlp.c_proj.modules_to_save` instead of the layer `base_model.model.transformer.h.3.mlp.c_proj.original_module`.

peft version : 0.3.0


Discussion:
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.


This doesn't seem solved yet from what I see in `ModulesToSaveWrapper`. When the adapter is disable, we probably still use `original_module` instead of `modules_to_save`.
Hi @glerzing 
Thanks for the issue, can you share a handy reproducible small snippet that best describes your problem?

Here it is :

``` Python
from peft import LoraConfig, get_peft_model, TaskType
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

peft_config = LoraConfig(
    task_type=TaskType.CAUSAL_LM,
    r=8,
    lora_alpha=32,
    lora_dropout=0.0,
    modules_to_save=["base_model.model.transformer.h.3.mlp"],
)

tokenizer = AutoTokenizer.from_pretrained("roneneldan/TinyStories-1M")
model = AutoModelForCausalLM.from_pretrained("roneneldan/TinyStories-1M")
model = get_peft_model(model, peft_config)

inputs = tokenizer("Hello world", return_tensors="pt")

initial_logits = model(**inputs, return_dict=True).logits

# Cause whatever loss to be backpropagated. The aim is just to change the parameter values.
loss = torch.nn.functional.binary_cross_entropy_with_logits(initial_logits[0][-1][:1], torch.tensor([0.53]))
loss.backward()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
optimizer.step()

# Check that the backpropagation worked
new_logits = model(**inputs, return_dict=True).logits
assert not torch.equal(initial_logits, new_logits)

model.disable_adapter_layers()
logits_without_adapter = model(**inputs, return_dict=True).logits
# Should trigger an error if the problem is not solved
assert torch.equal(initial_logits, logits_without_adapter)
```
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.


I could reproduce the issue, but commenting out this line fixed it:

`modules_to_save=["base_model.model.transformer.h.3.mlp"]`

The problem is that this layer will be updated during training in the usual fashion, i.e. without LoRA. Therefore, disabling the LoRA adapter will not revert the changes to that specific layer, which explains the different results.
The aim of my test is indeed to verify the correct behavior of  `disable_adapter_layers` on a model with `modules_to_save`.

Isn't the point of having duplicated the layers into `modules_to_save` and `original_module` to be able to disable the adapter with `disable_adapter_layers` and re-enable it ? It's practical, because it means that you can use both the original model and the fine-tuned model in one single model, without doubling the memory requirement. Only the fine-tuned, non-Lora layers are duplicated.
Yes, I think you're correct. I'll be adding a fix for this, thanks for digging into it.

Resolving patch:
diff --git a/src/peft/tuners/ia3.py b/src/peft/tuners/ia3.py
--- a/src/peft/tuners/ia3.py
+++ b/src/peft/tuners/ia3.py
@@ -303,6 +303,8 @@ def _set_adapter_layers(self, enabled=True):
         for module in self.model.modules():
             if isinstance(module, IA3Layer):
                 module.disable_adapters = False if enabled else True
+            elif isinstance(module, ModulesToSaveWrapper):
+                module.disable_adapters = False if enabled else True
 
     def enable_adapter_layers(self):
         self._set_adapter_layers(enabled=True)
diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py
--- a/src/peft/tuners/lora.py
+++ b/src/peft/tuners/lora.py
@@ -398,6 +398,8 @@ def _set_adapter_layers(self, enabled=True):
         for module in self.model.modules():
             if isinstance(module, LoraLayer):
                 module.disable_adapters = False if enabled else True
+            elif isinstance(module, ModulesToSaveWrapper):
+                module.disable_adapters = False if enabled else True
 
     def enable_adapter_layers(self):
         self._set_adapter_layers(enabled=True)
diff --git a/src/peft/utils/other.py b/src/peft/utils/other.py
--- a/src/peft/utils/other.py
+++ b/src/peft/utils/other.py
@@ -135,12 +135,13 @@ def __init__(self, module_to_save, adapter_name):
         self.modules_to_save = torch.nn.ModuleDict({})
         self.update(adapter_name)
         self.active_adapter = adapter_name
+        self.disable_adapters = False
 
     def update(self, adapter_name):
         self.modules_to_save.update(torch.nn.ModuleDict({adapter_name: copy.deepcopy(self.original_module)}))
 
     def forward(self, *args, **kwargs):
-        if self.active_adapter not in self.modules_to_save:
+        if self.disable_adapters or (self.active_adapter not in self.modules_to_save):
             return self.original_module(*args, **kwargs)
         return self.modules_to_save[self.active_adapter](*args, **kwargs)
 


Test patch:
diff --git a/tests/test_custom_models.py b/tests/test_custom_models.py
--- a/tests/test_custom_models.py
+++ b/tests/test_custom_models.py
@@ -35,6 +35,17 @@
     ("Vanilla MLP 2", "MLP", LoraConfig, {"target_modules": ["lin0"]}),
     ("Vanilla MLP 3", "MLP", LoraConfig, {"target_modules": ["lin1"]}),
     ("Vanilla MLP 4", "MLP", LoraConfig, {"target_modules": ["lin0", "lin1"]}),
+    ("Vanilla MLP 5", "MLP", LoraConfig, {"target_modules": ["lin0"], "modules_to_save": ["lin1"]}),
+    (
+        "Vanilla MLP 6",
+        "MLP",
+        LoraConfig,
+        {
+            "target_modules": ["lin0"],
+            "lora_alpha": 4,
+            "lora_dropout": 0.1,
+        },
+    ),
     ("Embedding + transformers Conv1D 1", "EmbConv1D", LoraConfig, {"target_modules": ["conv1d"]}),
     ("Embedding + transformers Conv1D 2", "EmbConv1D", LoraConfig, {"target_modules": ["emb"]}),
     ("Embedding + transformers Conv1D 3", "EmbConv1D", LoraConfig, {"target_modules": ["emb", "conv1d"]}),
@@ -227,7 +238,8 @@ def test_only_params_are_updated(self, test_name, model_id, config_cls, config_k
         self.assertEqual(params_before.keys(), params_after.keys())
         for name, param_before in params_before.items():
             param_after = params_after[name]
-            if "lora_" in name:
+            if ("lora_" in name) or ("modules_to_save" in name):
+                # target_modules and modules_to_save _are_ updated
                 self.assertFalse(torch.allclose(param_before, param_after, atol=tol, rtol=tol))
             else:
                 self.assertTrue(torch.allclose(param_before, param_after, atol=tol, rtol=tol))
@@ -262,8 +274,12 @@ def test_disable_adapters(self, test_name, model_id, config_cls, config_kwargs):
         with model.disable_adapter():
             outputs_disabled = model(**X)
 
+        # check that after leaving the disable_adapter context, everything is enabled again
+        outputs_enabled_after_disable = model(**X)
+
         self.assertFalse(torch.allclose(outputs_before, outputs_after))
         self.assertTrue(torch.allclose(outputs_before, outputs_disabled))
+        self.assertTrue(torch.allclose(outputs_after, outputs_enabled_after_disable))
 
     @parameterized.expand(TEST_CASES)
     def test_disable_adapter_with_bias_warns(self, test_name, model_id, config_cls, config_kwargs):
diff --git a/tests/testing_common.py b/tests/testing_common.py
--- a/tests/testing_common.py
+++ b/tests/testing_common.py
@@ -467,7 +467,7 @@ def _test_training(self, model_id, config_cls, config_kwargs):
         loss.backward()
         parameter_prefix = "ia3" if config_cls == IA3Config else "lora"
         for n, param in model.named_parameters():
-            if parameter_prefix in n:
+            if (parameter_prefix in n) or ("modules_to_save" in n):
                 self.assertIsNotNone(param.grad)
             else:
                 self.assertIsNone(param.grad)