FIX Reduce false positive missing keys when loading adapter (#2084)

When loading a PEFT adapter, a lot of missing keys are reported, because the
base model weights are not loaded. However, this is totally fine. Therefore,
those missing keys can be safely ignored.

When using from_pretrrained, the missing keys won't be returned to the user,
thus there is no room for confusion. But when using load_adapter, the missing
keys (and unexpected keys) are returned and can cause confusion. With this PR,
the missing keys are filtered to remove keys that are unrelated to the adapter.

A small gap is VB-LoRA which reports missing keys because the vector bank
parameters are actually only loaded once and then shared.

src/peft/peft_model.py

@@ -38,7 +38,7 @@
 from transformers.modeling_outputs import QuestionAnsweringModelOutput, SequenceClassifierOutput, TokenClassifierOutput
 from transformers.utils import PushToHubMixin
 
-from peft.utils.constants import DUMMY_MODEL_CONFIG
+from peft.utils.constants import DUMMY_MODEL_CONFIG, PEFT_TYPE_TO_PREFIX_MAPPING
 
 from . import __version__
 from .config import PeftConfig
@@ -1185,6 +1185,19 @@ def load_adapter(
  ignore_mismatched_sizes=ignore_mismatched_sizes,
  low_cpu_mem_usage=low_cpu_mem_usage,
  )
+
+ tuner = self.peft_config[adapter_name].peft_type
+ tuner_prefix = PEFT_TYPE_TO_PREFIX_MAPPING.get(tuner, "")
+ adapter_missing_keys = []
+
+ # Filter missing keys specific to the current adapter and tuner prefix.
+ for key in load_result.missing_keys:
+ if tuner_prefix in key and adapter_name in key:
+ adapter_missing_keys.append(key)
+
+ load_result.missing_keys.clear()
+ load_result.missing_keys.extend(adapter_missing_keys)
+
  if (
  (getattr(self, "hf_device_map", None) is not None)
  and (len(set(self.hf_device_map.values()).intersection({"cpu", "disk"})) > 0)

src/peft/utils/constants.py

@@ -15,6 +15,8 @@
 import torch
 from transformers import BloomPreTrainedModel
 
+from .peft_types import PeftType
+
 
 # needed for prefix-tuning of bloom model
 def bloom_model_postprocess_past_key_value(past_key_values):
@@ -284,6 +286,22 @@ def starcoder_model_postprocess_past_key_value(past_key_values):
  "qwen2": ["q_proj", "v_proj"],
 }
 
+PEFT_TYPE_TO_PREFIX_MAPPING = {
+ PeftType.IA3: "ia3_",
+ PeftType.LORA: "lora_",
+ PeftType.ADALORA: "lora_",
+ PeftType.LOHA: "hada_",
+ PeftType.LOKR: "lokr_",
+ PeftType.OFT: "oft_",
+ PeftType.POLY: "poly_",
+ PeftType.BOFT: "boft_",
+ PeftType.LN_TUNING: "ln_tuning_",
+ PeftType.VERA: "vera_lambda_",
+ PeftType.FOURIERFT: "fourierft_",
+ PeftType.HRA: "hra_",
+ PeftType.VBLORA: "vblora_",
+}
+
 WEIGHTS_NAME = "adapter_model.bin"
 SAFETENSORS_WEIGHTS_NAME = "adapter_model.safetensors"
 CONFIG_NAME = "adapter_config.json"

src/peft/utils/save_and_load.py

@@ -24,6 +24,7 @@
 from packaging import version
 from safetensors.torch import load_file as safe_load_file
 
+from .constants import PEFT_TYPE_TO_PREFIX_MAPPING
 from .other import (
  EMBEDDING_LAYER_NAMES,
  SAFETENSORS_WEIGHTS_NAME,
@@ -357,21 +358,7 @@ def set_peft_model_state_dict(
  PeftType.VBLORA,
  ):
  peft_model_state_dict = {}
- parameter_prefix = {
- PeftType.IA3: "ia3_",
- PeftType.LORA: "lora_",
- PeftType.ADALORA: "lora_",
- PeftType.LOHA: "hada_",
- PeftType.LOKR: "lokr_",
- PeftType.OFT: "oft_",
- PeftType.POLY: "poly_",
- PeftType.BOFT: "boft_",
- PeftType.LN_TUNING: "ln_tuning_",
- PeftType.VERA: "vera_lambda_",
- PeftType.FOURIERFT: "fourierft_",
- PeftType.HRA: "hra_",
- PeftType.VBLORA: "vblora_",
- }[config.peft_type]
+ parameter_prefix = PEFT_TYPE_TO_PREFIX_MAPPING[config.peft_type]
  if config.peft_type == PeftType.VBLORA and config.save_only_topk_weights:
  num_vectors, _ = model.vblora_vector_bank[adapter_name].shape
  state_dict_keys = list(state_dict.keys())

tests/testing_common.py

@@ -532,8 +532,16 @@ def _test_load_multiple_adapters(self, model_id, config_cls, config_kwargs):
 
  model = self.transformers_class.from_pretrained(model_id).to(self.torch_device)
  model = PeftModel.from_pretrained(model, tmp_dirname, torch_device=self.torch_device)
- model.load_adapter(tmp_dirname, adapter_name="other")
- model.load_adapter(tmp_dirname, adapter_name="yet-another")
+
+ load_result1 = model.load_adapter(tmp_dirname, adapter_name="other")
+ load_result2 = model.load_adapter(tmp_dirname, adapter_name="yet-another")
+
+ # VBLoRA uses a shared "vblora_vector_bank" across all layers, causing it to appear
+ # in the missing keys list, which leads to failed test cases. So
+ # skipping the missing keys check for VBLoRA.
+ if config.peft_type != "VBLORA":
+ assert load_result1.missing_keys == []
+ assert load_result2.missing_keys == []
 
  def _test_merge_layers_fp16(self, model_id, config_cls, config_kwargs):
  if config_cls not in (LoraConfig, IA3Config, AdaLoraConfig, LoHaConfig, LoKrConfig, VBLoRAConfig):