Fix not detecting regex-targeted embedding layer (#2649)

This issue was found in PR #2638 and is defined thusly:

> When calling `get_peft_model_state_dict(..., save_embedding_layers="auto")` we check if the
> embedding layer is targetted to determine if the embedding layers need saving. This is not
> done when `PeftConfig.target_modules` is a regex-string, potentially missing to save embeddings.

This is fixed by adding a check similar to the existing query of whether `EMBEDDING_LAYER_NAMES` is
a subset of the defined target modules, only that the regex matching from `BaseTuner.inject_adapter`
is used. To avoid code duplication, the matching was moved to its own utility function
`match_target_against_key`.

The main complication was to define the test-cases as it was non-trivial to find what the meaning
of `save_embedding_layers="auto"` entails. I've assembled a list of cases that I think are correct
in the corresponding unit test.

Author: githubnemo

src/peft/tuners/tuners_utils.py

@@ -38,7 +38,7 @@
     SEQ_CLS_HEAD_NAMES,
 )
 from peft.utils.integrations import init_empty_weights
-from peft.utils.other import AuxiliaryTrainingWrapper, set_additional_trainable_modules
+from peft.utils.other import AuxiliaryTrainingWrapper, match_target_against_key, set_additional_trainable_modules
 from peft.utils.peft_types import PeftType, TaskType
 
 from ..config import PeftConfig
@@ -1133,7 +1133,7 @@ def check_target_module_exists(config, key: str) -> bool | re.Match[str] | None:
         return False
 
     if isinstance(config.target_modules, str):
-        target_module_found = re.fullmatch(config.target_modules, key)
+        target_module_found = match_target_against_key(config.target_modules, key)
     elif key in config.target_modules:
         # this module is specified directly in target_modules
         target_module_found = True

src/peft/utils/other.py

@@ -1229,6 +1229,14 @@ def check_file_exists_on_hf_hub(repo_id: str, filename: str, **kwargs) -> Option
     return exists
 
 
+def match_target_against_key(target_pattern: str, key: str):
+    """Backing function for `target_modules` config parameter.
+
+    Having this as its own function ensures that target key matching can be implemented in the same way everywhere.
+    """
+    return re.fullmatch(target_pattern, key)
+
+
 def get_pattern_key(pattern_keys: Sequence[str], key_to_match: str) -> str:
     """Match a substring of key_to_match in pattern keys"""
     for key in pattern_keys:

src/peft/utils/save_and_load.py

@@ -28,13 +28,15 @@
 
 from peft.mapping import PEFT_TYPE_TO_PREFIX_MAPPING
 
+from .constants import INCLUDE_LINEAR_LAYERS_SHORTHAND
 from .other import (
     EMBEDDING_LAYER_NAMES,
     SAFETENSORS_WEIGHTS_NAME,
     WEIGHTS_NAME,
     AuxiliaryTrainingWrapper,
     check_file_exists_on_hf_hub,
     infer_device,
+    match_target_against_key,
 )
 from .peft_types import PeftType
 
@@ -235,23 +237,35 @@ def renamed_dora_weights(k):
             )
 
     # DEAL WITH EMBEDDINGS
-    # check the common embedding layers in `target_modules` to reset `save_embedding_layers` if necessary
-    is_embedding_in_target_modules = False
+    #
+    # save_embedding_layer="auto" needs to check the following logic:
+    #
+    # - when vocab size was NOT changed, embeddings should be saved only when targeted
+    # but not when
+    # - using PeftType.TRAINABLE_TOKENS
+    # - LoRA using trainable_token_indices (since their goal is to space-efficient)
+    # but
+    # - when vocab size was changed, embeddings should be saved automatically regardless to cover this
+    #   scenario: 1) fine-tune embedding, 2) resize embedding, 3) train with trainable tokens
+    #
     embedding_is_targeted = False
     if hasattr(config, "target_modules"):
-        if isinstance(config.target_modules, str):
-            # TODO: implement this; note: this change is not directly related to the PR, the bug already existed b4
-            pass
+        if isinstance(config.target_modules, str) and config.target_modules != INCLUDE_LINEAR_LAYERS_SHORTHAND:
+            embedding_is_targeted = any(
+                match_target_against_key(config.target_modules, k)
+                for k, _ in model.get_base_model().named_modules()
+                if any(re.match(rf"(.*\.)?{e}$", k) for e in EMBEDDING_LAYER_NAMES)
+            )
         elif config.target_modules:
             embedding_is_targeted = any(k in config.target_modules for k in EMBEDDING_LAYER_NAMES)
-    if (
-        save_embedding_layers == "auto"
-        and hasattr(config, "target_modules")
-        and embedding_is_targeted
-        and config.peft_type != PeftType.TRAINABLE_TOKENS
-    ):
+
+    using_trainable_tokens = (
+        config.peft_type == PeftType.TRAINABLE_TOKENS or getattr(config, "trainable_token_indices", None) is not None
+    )
+
+    if save_embedding_layers == "auto" and embedding_is_targeted and not using_trainable_tokens:
         warnings.warn("Setting `save_embedding_layers` to `True` as embedding layers found in `target_modules`.")
-        save_embedding_layers = is_embedding_in_target_modules = True
+        save_embedding_layers = True
     elif save_embedding_layers == "auto":
         vocab_size = getattr(getattr(model, "config", None), "vocab_size", None)
         model_id = getattr(config, "base_model_name_or_path", None)
@@ -289,9 +303,10 @@ def renamed_dora_weights(k):
 
     if save_embedding_layers and hasattr(model, "get_input_embeddings"):
         for layer in [model.get_input_embeddings(), model.get_output_embeddings()]:
-            if not is_embedding_in_target_modules or has_valid_embedding_base_layer(layer):
-                # support from version >= 0.6.2
-                embedding_module_name = get_embedding_layer_name(model, layer, is_embedding_in_target_modules)
+            # Either the layer is not targeted, then it must have been resized and needs saving. Or it is targeted and
+            # therefore has a valid base layer, then we'll save it as well.
+            if not embedding_is_targeted or has_valid_embedding_base_layer(layer):
+                embedding_module_name = get_embedding_layer_name(model, layer, embedding_is_targeted)
                 if embedding_module_name:
                     to_return.update({k: v for k, v in state_dict.items() if embedding_module_name in k})
     elif save_embedding_layers:

tests/test_custom_models.py

@@ -2392,52 +2392,62 @@ def test_non_existing_model_card(self):
         assert len(model_card) > 1000
 
     @pytest.mark.parametrize("save_embedding_layers", ["auto", True, False])
-    def test_targeting_lora_to_embedding_layer(self, save_embedding_layers):
+    @pytest.mark.parametrize(
+        "peft_config",
+        [
+            (LoraConfig(target_modules=["lin0", "embed_tokens"], init_lora_weights=False)),
+            (LoraConfig(target_modules=r"^embed_tokens", init_lora_weights=False)),
+        ],
+    )
+    def test_save_pretrained_targeting_lora_to_embedding_layer(self, save_embedding_layers, tmp_path, peft_config):
         model = ModelEmbWithEmbeddingUtils()
-        config = LoraConfig(target_modules=["embed_tokens", "lin0"], init_lora_weights=False)
-        model = get_peft_model(model, config)
+        model = get_peft_model(model, peft_config)
 
-        with tempfile.TemporaryDirectory() as tmp_dirname:
-            if save_embedding_layers == "auto":
-                # assert warning
-                msg_start = "Setting `save_embedding_layers` to `True` as embedding layers found in `target_modules`."
-                with pytest.warns(UserWarning, match=msg_start):
-                    model.save_pretrained(tmp_dirname, save_embedding_layers=save_embedding_layers)
-            else:
-                model.save_pretrained(tmp_dirname, save_embedding_layers=save_embedding_layers)
-            from safetensors.torch import load_file as safe_load_file
-
-            state_dict = safe_load_file(os.path.join(tmp_dirname, "adapter_model.safetensors"))
-            if save_embedding_layers in ["auto", True]:
-                assert "base_model.model.embed_tokens.base_layer.weight" in state_dict
-                assert torch.allclose(
-                    model.base_model.model.embed_tokens.base_layer.weight,
-                    state_dict["base_model.model.embed_tokens.base_layer.weight"],
-                )
-            else:
-                assert "base_model.model.embed_tokens.base_layer.weight" not in state_dict
-            del state_dict
+        if save_embedding_layers == "auto":
+            # assert warning
+            msg_start = "Setting `save_embedding_layers` to `True` as embedding layers found in `target_modules`."
+            with pytest.warns(UserWarning, match=msg_start):
+                model.save_pretrained(tmp_path, save_embedding_layers=save_embedding_layers)
+        else:
+            model.save_pretrained(tmp_path, save_embedding_layers=save_embedding_layers)
+
+        state_dict = safe_load_file(tmp_path / "adapter_model.safetensors")
+        contains_embedding = "base_model.model.embed_tokens.base_layer.weight" in state_dict
+
+        if save_embedding_layers in ["auto", True]:
+            assert contains_embedding
+            assert torch.allclose(
+                model.base_model.model.embed_tokens.base_layer.weight,
+                state_dict["base_model.model.embed_tokens.base_layer.weight"],
+            )
+        else:
+            assert not contains_embedding
 
     @pytest.mark.parametrize("save_embedding_layers", ["auto", True, False])
-    def test_targeting_lora_to_embedding_layer_non_transformers(self, save_embedding_layers):
+    @pytest.mark.parametrize(
+        "peft_config",
+        [
+            (LoraConfig(target_modules=["lin0", "emb"], init_lora_weights=False)),
+            (LoraConfig(target_modules=r"^emb", init_lora_weights=False)),
+        ],
+    )
+    def test_save_pretrained_targeting_lora_to_embedding_layer_non_transformers(
+        self, save_embedding_layers, tmp_path, peft_config
+    ):
         model = ModelEmbConv1D()
-        config = LoraConfig(target_modules=["emb", "lin0"], init_lora_weights=False)
-        model = get_peft_model(model, config)
-
-        with tempfile.TemporaryDirectory() as tmp_dirname:
-            if save_embedding_layers is True:
-                with pytest.warns(
-                    UserWarning,
-                    match=r"Could not identify embedding layer\(s\) because the model is not a 🤗 transformers model\.",
-                ):
-                    model.save_pretrained(tmp_dirname, save_embedding_layers=save_embedding_layers)
-            else:
-                model.save_pretrained(tmp_dirname, save_embedding_layers=save_embedding_layers)
-            from safetensors.torch import load_file as safe_load_file
+        model = get_peft_model(model, peft_config)
+
+        if save_embedding_layers is True:
+            with pytest.warns(
+                UserWarning,
+                match=r"Could not identify embedding layer\(s\) because the model is not a 🤗 transformers model\.",
+            ):
+                model.save_pretrained(tmp_path, save_embedding_layers=save_embedding_layers)
+        else:
+            model.save_pretrained(tmp_path, save_embedding_layers=save_embedding_layers)
 
-            state_dict = safe_load_file(os.path.join(tmp_dirname, "adapter_model.safetensors"))
-            assert "base_model.model.emb.base_layer.weight" not in state_dict
-            del state_dict
+        state_dict = safe_load_file(tmp_path / "adapter_model.safetensors")
+        assert "base_model.model.emb.base_layer.weight" not in state_dict
 
     def test_load_resized_embedding_ignore_mismatched_sizes(self):
         # issue #1605

tests/test_decoder_models.py

@@ -17,6 +17,7 @@
 
 import pytest
 import torch
+from safetensors.torch import load_file as safe_load_file
 from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
@@ -46,7 +47,7 @@
     get_peft_model,
 )
 
-from .testing_common import PeftCommonTester
+from .testing_common import PeftCommonTester, hub_online_once
 from .testing_utils import device_count, load_dataset_english_quotes, set_init_weights_false
 
 
@@ -680,3 +681,38 @@ def process(samples):
                 data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
             )
             trainer.train()
+
+    @pytest.mark.parametrize("save_embedding_layers", ["auto", True, False])
+    @pytest.mark.parametrize(
+        "peft_config",
+        [
+            (LoraConfig(target_modules=["lin0", "embed_tokens"], init_lora_weights=False)),
+            (LoraConfig(target_modules=r".*\.embed_tokens", init_lora_weights=False)),
+        ],
+    )
+    def test_save_pretrained_targeting_lora_to_embedding_layer(self, save_embedding_layers, tmp_path, peft_config):
+        model_id = "trl-internal-testing/tiny-random-LlamaForCausalLM"
+
+        with hub_online_once(model_id):
+            model = AutoModelForCausalLM.from_pretrained(model_id)
+            model = get_peft_model(model, peft_config)
+
+            if save_embedding_layers == "auto":
+                # assert warning
+                msg_start = "Setting `save_embedding_layers` to `True` as embedding layers found in `target_modules`."
+                with pytest.warns(UserWarning, match=msg_start):
+                    model.save_pretrained(tmp_path, save_embedding_layers=save_embedding_layers)
+            else:
+                model.save_pretrained(tmp_path, save_embedding_layers=save_embedding_layers)
+
+            state_dict = safe_load_file(tmp_path / "adapter_model.safetensors")
+            contains_embedding = "base_model.model.model.embed_tokens.base_layer.weight" in state_dict
+
+            if save_embedding_layers in ["auto", True]:
+                assert contains_embedding
+                assert torch.allclose(
+                    model.base_model.model.model.embed_tokens.base_layer.weight,
+                    state_dict["base_model.model.model.embed_tokens.base_layer.weight"],
+                )
+            else:
+                assert not contains_embedding

tests/test_trainable_tokens.py

@@ -18,6 +18,7 @@
 
 import pytest
 import torch
+from safetensors.torch import load_file as safe_load_file
 from transformers import AutoModelForCausalLM, AutoModelForSeq2SeqLM, AutoTokenizer
 
 from peft import AutoPeftModel, LoraConfig, PeftModel, TrainableTokensConfig, get_peft_model
@@ -885,3 +886,31 @@ def test_embedding_name_is_used_when_given_combined(self, model_embed_multiple):
 
         assert isinstance(peft_model.model.embed_in_2, TrainableTokensWrapper)
         assert not isinstance(peft_model.model.embed_in, TrainableTokensWrapper)
+
+    @pytest.mark.parametrize("resize_embedding", [True, False])
+    @pytest.mark.parametrize(
+        "peft_config",
+        [
+            LoraConfig(target_modules="all-linear", trainable_token_indices=[1, 2, 3]),
+            TrainableTokensConfig(target_modules=None, token_indices=[1, 2, 3]),
+        ],
+    )
+    def test_save_pretrained_auto(self, model, resize_embedding, peft_config, tmp_path):
+        # make sure that embeddings are saved alongside trainable token weights but only when
+        # the we detect the embedding to be resized (as detected by save_embedding_layers="auto")
+        if resize_embedding:
+            model.resize_token_embeddings(model.config.vocab_size + 2)
+        peft_model = get_peft_model(model, peft_config)
+
+        peft_model.save_pretrained(tmp_path, save_embedding_layers="auto")
+        state_dict = safe_load_file(tmp_path / "adapter_model.safetensors")
+
+        if isinstance(peft_config, TrainableTokensConfig):
+            contains_embedding = "base_model.model.model.embed_tokens.base_layer.weight" in state_dict
+        else:
+            contains_embedding = "base_model.model.model.embed_tokens.token_adapter.base_layer.weight" in state_dict
+
+        if resize_embedding:
+            assert contains_embedding
+        else:
+            assert not contains_embedding