Repository: huggingface/peft

Issue:
'Embedding' object has no attribute 'device' for MPT
I followed the test process of MPT notebook, but encountered an error during target training: 'Embedding' object has no attribute 'device'. I want to ask how to solve this.


Discussion:
Detail:
<img width="951" alt="bug" src="https://github.com/huggingface/peft/assets/52874161/02e2577d-d652-4f0a-a796-dcaa072c3fc5">

I believe the bug is in 'self.word_embeddings = transformer_backbone.get_submodule(named_param.replace(".weight", ""))' in peft_model, there don't have 'device',   'self.word_embeddings.device = self.device' may need to be added.
I can reproduce the error. The solution in PEFT would be to change this line:

https://github.com/huggingface/peft/blob/cfe35a7878b44e017836f0b0c0c3b3e9e0cb738b/src/peft/tuners/multitask_prompt_tuning/model.py#L72

to `map_location=word_embeddings.weight.device`.
@BenjaminBossan Thanks！

Let's keep this issue open, as we should fix this issue in PEFT :)
> Let's keep this issue open, as we should fix this issue in PEFT :)
ok👌

Resolving patch:
diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py
--- a/src/peft/peft_model.py
+++ b/src/peft/peft_model.py
@@ -1153,7 +1153,7 @@ def generate(self, *args, **kwargs):
             self.base_model.prepare_inputs_for_generation = self.base_model_prepare_inputs_for_generation
             return outputs
 
-    def prepare_inputs_for_generation(self, *args, task_ids: torch.Tensor = None, **kwargs):
+    def prepare_inputs_for_generation(self, *args, task_ids: Optional[torch.Tensor] = None, **kwargs):
         peft_config = self.active_peft_config
         model_kwargs = self.base_model_prepare_inputs_for_generation(*args, **kwargs)
 
diff --git a/src/peft/tuners/multitask_prompt_tuning/model.py b/src/peft/tuners/multitask_prompt_tuning/model.py
--- a/src/peft/tuners/multitask_prompt_tuning/model.py
+++ b/src/peft/tuners/multitask_prompt_tuning/model.py
@@ -66,9 +66,10 @@ def __init__(self, config: MultitaskPromptTuningConfig, word_embeddings):
                     "init method"
                 )
 
+            # TODO: There should be an option for safetensors
             state_dict: dict = torch.load(
                 config.prompt_tuning_init_state_dict_path,
-                map_location=word_embeddings.device,
+                map_location=word_embeddings.weight.device,
             )
 
         if config.prompt_tuning_init in [


Test patch:
diff --git a/tests/test_multitask_prompt_tuning.py b/tests/test_multitask_prompt_tuning.py
--- a/tests/test_multitask_prompt_tuning.py
+++ b/tests/test_multitask_prompt_tuning.py
@@ -17,13 +17,15 @@
 import tempfile
 from unittest import TestCase
 
+import pytest
 import torch
+from parameterized import parameterized
 from torch.testing import assert_close
 
 from peft.mapping import get_peft_model
 from peft.peft_model import PeftModel
-from peft.tuners.multitask_prompt_tuning import MultitaskPromptTuningConfig
-from peft.utils.other import prepare_model_for_int8_training
+from peft.tuners.multitask_prompt_tuning import MultitaskPromptTuningConfig, MultitaskPromptTuningInit
+from peft.utils.other import WEIGHTS_NAME, prepare_model_for_int8_training
 from peft.utils.save_and_load import get_peft_model_state_dict
 from tests.testing_common import PeftCommonTester
 
@@ -73,7 +75,9 @@ def _create_multitask_prompt_tuning_config(cls) -> MultitaskPromptTuningConfig:
             task_type="CAUSAL_LM",
             num_virtual_tokens=50,
             num_tasks=3,
-            prompt_tuning_init_text="classify the following into either positive or negative, or entailment, neutral or contradiction:",
+            prompt_tuning_init_text=(
+                "classify the following into either positive or negative, or entailment, neutral or contradiction:"
+            ),
         )
 
     def test_prepare_for_training(self) -> None:
@@ -240,3 +244,61 @@ def test_bf16_inference(self) -> None:
         mpt = get_peft_model(original, self._create_multitask_prompt_tuning_config())
         mpt = mpt.to(self.torch_device)
         _ = mpt.generate(input_ids=input_ids, task_ids=task_ids)
+
+    def test_generate_text_with_random_init(self) -> None:
+        model = LlamaForCausalLM(self._create_test_llama_config())
+
+        config = self._create_multitask_prompt_tuning_config()
+        config.prompt_tuning_init = MultitaskPromptTuningInit.RANDOM
+
+        model = get_peft_model(model, config)
+        model = model.to(self.torch_device)
+
+        input_ids = torch.LongTensor([[1, 1, 1], [2, 1, 2]]).to(self.torch_device)
+        attention_mask = torch.LongTensor([[1, 1, 1], [1, 0, 1]]).to(self.torch_device)
+        task_ids = torch.LongTensor([0]).to(self.torch_device)
+
+        # check if `generate` works
+        _ = model.generate(input_ids=input_ids, attention_mask=attention_mask, task_ids=task_ids)
+
+        with pytest.raises(ValueError):
+            # check if `generate` raises an error if task_ids are not passed
+            _ = model.generate(input_ids, attention_mask=attention_mask)
+
+    @parameterized.expand(
+        [
+            MultitaskPromptTuningInit.AVERAGE_SOURCE_TASKS,
+            MultitaskPromptTuningInit.EXACT_SOURCE_TASK,
+            MultitaskPromptTuningInit.ONLY_SOURCE_SHARED,
+        ],
+    )
+    def test_generate_text_with_other_init(self, prompt_tuning_init) -> None:
+        with tempfile.TemporaryDirectory() as tmp_dirname:
+            model = LlamaForCausalLM(self._create_test_llama_config())
+            model = get_peft_model(model, self._create_multitask_prompt_tuning_config())
+            model.save_pretrained(tmp_dirname, safe_serialization=False)  # bc torch.load is used
+
+            config = MultitaskPromptTuningConfig(
+                task_type="CAUSAL_LM",
+                num_virtual_tokens=50,
+                num_tasks=1,
+                prompt_tuning_init_text=(
+                    "classify the following into either positive or negative, or entailment, neutral or contradiction:"
+                ),
+                prompt_tuning_init=prompt_tuning_init,
+                prompt_tuning_init_state_dict_path=os.path.join(tmp_dirname, WEIGHTS_NAME),
+            )
+            model = LlamaForCausalLM(self._create_test_llama_config())
+            model = get_peft_model(model, config)
+            model = model.to(self.torch_device)
+
+            input_ids = torch.LongTensor([[1, 1, 1], [2, 1, 2]]).to(self.torch_device)
+            attention_mask = torch.LongTensor([[1, 1, 1], [1, 0, 1]]).to(self.torch_device)
+            task_ids = torch.LongTensor([0]).to(self.torch_device)
+
+            # check if `generate` works
+            _ = model.generate(input_ids=input_ids, attention_mask=attention_mask, task_ids=task_ids)
+
+            with pytest.raises(ValueError):
+                # check if `generate` raises an error if task_ids are not passed
+                _ = model.generate(input_ids, attention_mask=attention_mask)