Repository: huggingface/peft

Issue:
merge_and_unload issue?
### System Info

I am using the latest dev version of transformers, accelerate and peft (installed via !pip install -q -U git+<URL>) installed via Google Colab. 

This worked a few days ago, but now when I attempt to merge an adapter back into the base model and then save to hub, the size is much smaller than that of the base model and it can't be loaded (generates error "Cannot copy out of meta tensor; no data!" when I attempt to do so). 

The function I am using to merge the PeftModel back into the base model is:
**(Code begins here)**

```python
def merge_adapter(lora_id, merged_id):
    config = PeftConfig.from_pretrained(lora_id)

    model_id = config.base_model_name_or_path
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        torch_dtype=dtype,
        device_map="auto",
        offload_folder="offload"
        )

    adapter = PeftModel.from_pretrained(
        model,
        lora_id,
        torch_dtype=dtype,
        device_map="auto",
        offload_folder="offload"
    )

    model = adapter.merge_and_unload(progressbar=True)
    tokenizer = AutoTokenizer.from_pretrained(model_id) #, use_fast=False)

    model.save_pretrained(
        merged_id,
        push_to_hub=True,
        repo_id=merged_id,
        private=True,
        #max_shard_size="4GB"
    #,
    )

    tokenizer.save_pretrained(
        merged_id,
        push_to_hub=True,
        repo_id=merged_id,
        private=True,
    )
```

The base model is 'meta-llama/Llama-2-7b-hf', which has two bin files, 1 @ 9.98GB and 1 @ 3.5GB. Previously when I ran this code, the merged model would be the same size. Now it only produces a single file @ 1.07GB. 

This may be an error with the library, although I am not seeing any bug reports to indicate this. It may also be an error with the training code in the first place, my upload code, the HF library or anything else. 

If anyone has any solutions, please let me know. Otherwise, if this is a bug, I guess this is my first bug report. 


### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder
- [X] My own task or dataset (give details below)

### Reproduction

This is the Colab I have been using. The data and LoRAs are private, but the two I am playing with are LLama 2 7B QLoRA fine-tuned on 1) a chat dataset and 2) a big hunk of raw text split into paragraphs (a piece of fanfiction). 

https://colab.research.google.com/drive/15g2NU2wJ9fOvY3PJCCN5dVDYV8KSXbeS?usp=sharing

### Expected behavior

The merged model should be created and be the same size as the base model, and I should be able to load it using AutoModelForCausalLM.from_pretrained as I was able to a few days ago. 


Discussion:
i meet the same issue.
when i merge lora to llama model, and the save it .the model size is smaller than before
I think the issue is that the layers that are on "meta" device are not properly handled in this case. Their weights are not loaded and the lora weights are also just on meta device. The result you see probably only contains the weights that were actually loaded, missing all the meta device weights.

Honestly, I'm not completely sure how best to handle this situation. Maybe @pacman100 or @younesbelkada have a good idea.
@BenjaminBossan ,yes.
when i checked. i find that some weight load on cpu , some on gpu . and the cpu weight is the meta device weights.
how should i change the meta device to cpu or GPU ?
Sorry, I'm really not sure what a solution would look like here. The reason why the weights are on meta device is that they needed to be offloaded for lack of memory, so it's not as simple as just loading everything. Hopefully one of the others can shine some light on this.
This issue may occur when the GPU is occupied by other processes. device="auto" may not be able to load models into GPUs, so try to ensure that device is free or set device="cpu" instead.
I've encountered this issue, the model was trained and merge with the same GPU(3090 24GB), VRAM is sufficient for lora training and inferencing, but after merged i got this issue
> This issue may occur when the GPU is occupied by other processes. device="auto" may not be able to load models into GPUs, so try to ensure that device is free or set device="cpu" instead.

Setting `device_map="cpu"` from `auto` sovled the issue for me
I think that merge and unload is not supported for models that are offloaded into disk / cpu as accelerate puts the offloaded weights on the meta device. This is a bug we should look into
how to merge ptuning model?
Hi @xiaobai52HZ I think that sadly currently merging p-tuning models is not supported. cc @pacman100 @BenjaminBossan 
Indeed.
Can we add boolean arguments to infer device map to disable all off-loading? 
I still want the benefits of using multiple GPUs but raise an error if the weights cannot fit with GPUs alone

Resolving patch:
diff --git a/src/peft/tuners/lora/model.py b/src/peft/tuners/lora/model.py
--- a/src/peft/tuners/lora/model.py
+++ b/src/peft/tuners/lora/model.py
@@ -31,7 +31,7 @@
 from transformers.pytorch_utils import Conv1D
 
 from peft.import_utils import is_bnb_4bit_available, is_bnb_available
-from peft.tuners.tuners_utils import BaseTuner, BaseTunerLayer, check_target_module_exists
+from peft.tuners.tuners_utils import BaseTuner, BaseTunerLayer, check_target_module_exists, onload_layer
 from peft.utils import (
     TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING,
     ModulesToSaveWrapper,
@@ -437,7 +437,6 @@ def _unload_and_optionally_merge(
             if getattr(self.model, "quantization_method", None) == "gptq":
                 raise ValueError("Cannot merge LORA layers when the model is gptq quantized")
 
-        self._unloading_checks(adapter_names)
         key_list = [key for key, _ in self.model.named_modules() if self.prefix not in key]
         desc = "Unloading " + ("and merging " if merge else "") + "model"
         for key in tqdm(key_list, disable=not progressbar, desc=desc):
@@ -445,14 +444,14 @@ def _unload_and_optionally_merge(
                 parent, target, target_name = _get_submodules(self.model, key)
             except AttributeError:
                 continue
-
-            if hasattr(target, "base_layer"):
-                if merge:
-                    target.merge(safe_merge=safe_merge, adapter_names=adapter_names)
-                self._replace_module(parent, target_name, target.get_base_layer(), target)
-            elif isinstance(target, ModulesToSaveWrapper):
-                # save any additional trainable modules part of `modules_to_save`
-                setattr(parent, target_name, target.modules_to_save[target.active_adapter])
+            with onload_layer(target):
+                if hasattr(target, "base_layer"):
+                    if merge:
+                        target.merge(safe_merge=safe_merge, adapter_names=adapter_names)
+                    self._replace_module(parent, target_name, target.get_base_layer(), target)
+                elif isinstance(target, ModulesToSaveWrapper):
+                    # save any additional trainable modules part of `modules_to_save`
+                    setattr(parent, target_name, target.modules_to_save[target.active_adapter])
 
         return self.model
 
diff --git a/src/peft/tuners/tuners_utils.py b/src/peft/tuners/tuners_utils.py
--- a/src/peft/tuners/tuners_utils.py
+++ b/src/peft/tuners/tuners_utils.py
@@ -18,9 +18,12 @@
 import re
 import warnings
 from abc import ABC, abstractmethod
+from contextlib import contextmanager
 from typing import Any, List, Optional, Union
 
 import torch
+from accelerate.hooks import AlignDevicesHook
+from accelerate.utils import named_module_tensors, offload_state_dict
 from torch import nn
 
 from peft.utils import COMMON_LAYERS_PATTERN
@@ -32,6 +35,58 @@
 logger = logging.getLogger(__name__)
 
 
+@contextmanager
+def onload_layer(layer):
+    r"""
+    A utility for modifying a module containing one or more tuners and a base layer, any of which are offloaded to the
+    CPU or disk. Moves a module's sub-modules to the execution device before some action is performed, after that the
+    base layer state dictionary is re-assigned (if that layer was offloaded to the disk) and finally the parameters are
+    offloaded.
+
+    If the module has no offloaded sub-modules, this function does nothing.
+
+    Args:
+        layer ('torch.nn.Module'):
+            layer with tuners to be merged
+    """
+
+    offloaded_modules = []
+    for name, module in layer.named_modules():
+        if name in ["", "base_layer"]:
+            continue
+        if hasattr(module, "_hf_hook") and isinstance(module._hf_hook, AlignDevicesHook) and module._hf_hook.offload:
+            module._hf_hook.pre_forward(module)
+            offloaded_modules.append(module)
+
+    base_layer_offload = False
+    if hasattr(layer, "base_layer") and (
+        hasattr(layer.base_layer, "_hf_hook")
+        and isinstance(layer.base_layer._hf_hook, AlignDevicesHook)
+        and layer.base_layer._hf_hook.offload
+    ):
+        if torch.device("meta") in layer.base_layer._hf_hook.original_devices.values():
+            # retrieve the name of the original disk-offload directory
+            offload_folder = layer.base_layer._hf_hook.weights_map.dataset.save_folder
+        layer.base_layer._hf_hook.pre_forward(layer.base_layer)
+        base_layer_offload = True
+
+    yield
+
+    for module in offloaded_modules:
+        module._hf_hook.post_forward(module, torch.tensor([]))
+
+    if base_layer_offload:
+        # re-make weights map (must be on cpu to send params to the disk via memmap if disk offload)
+        layer.base_layer._hf_hook.weights_map = {
+            name: param.to("cpu") for name, param in named_module_tensors(layer.base_layer)
+        }
+        # offload weights map to disk if original device is the disk
+        if torch.device("meta") in layer.base_layer._hf_hook.original_devices.values():
+            # rewrite directory with merged weights
+            offload_state_dict(offload_folder, layer.base_layer._hf_hook.weights_map)
+        layer.base_layer._hf_hook.post_forward(layer.base_layer, torch.tensor([]))
+
+
 class BaseTuner(nn.Module, ABC):
     r"""
     A base tuner model that provides the common methods and attributes for all tuners that are injectable into a
@@ -284,7 +339,8 @@ def merge_adapter(self, adapter_names: Optional[list[str]] = None) -> None:
         """
         for module in self.model.modules():
             if isinstance(module, BaseTunerLayer):
-                module.merge(adapter_names=adapter_names)
+                with onload_layer(module):
+                    module.merge(adapter_names=adapter_names)
 
     def unmerge_adapter(self):
         """
@@ -292,7 +348,8 @@ def unmerge_adapter(self):
         """
         for module in self.model.modules():
             if isinstance(module, BaseTunerLayer):
-                module.unmerge()
+                with onload_layer(module):
+                    module.unmerge()
 
     def _unloading_checks(self, adapter_names: Optional[List[str]]):
         adapters_to_consider = adapter_names or self.active_adapters


Test patch:
diff --git a/tests/test_gpu_examples.py b/tests/test_gpu_examples.py
--- a/tests/test_gpu_examples.py
+++ b/tests/test_gpu_examples.py
@@ -21,6 +21,7 @@
 
 import pytest
 import torch
+from accelerate import infer_auto_device_map
 from accelerate.test_utils.testing import run_command
 from accelerate.utils import patch_environment
 from datasets import Audio, DatasetDict, load_dataset
@@ -43,6 +44,7 @@
     AdaLoraConfig,
     LoftQConfig,
     LoraConfig,
+    PeftModel,
     TaskType,
     get_peft_model,
     prepare_model_for_int8_training,
@@ -943,6 +945,65 @@ def test_causal_lm_training_multi_gpu(self):
             self.assertIsNotNone(trainer.state.log_history[-1]["train_loss"])
 
 
+@require_torch_gpu
+class OffloadSaveTests(unittest.TestCase):
+    def setUp(self):
+        self.causal_lm_model_id = "gpt2"
+
+    def tearDown(self):
+        r"""
+        Efficient mechanism to free GPU memory after each test. Based on
+        https://github.com/huggingface/transformers/issues/21094
+        """
+        gc.collect()
+        torch.cuda.empty_cache()
+
+    @pytest.mark.single_gpu_tests
+    @require_torch_gpu
+    def test_offload_merge(self):
+        r"""
+        Test merging, unmerging, and unloading of a model with CPU- offloaded modules.
+        """
+        torch.manual_seed(0)
+        model = AutoModelForCausalLM.from_pretrained(self.causal_lm_model_id)
+        tokenizer = AutoTokenizer.from_pretrained(self.causal_lm_model_id)
+        # TODO: add disk offload once PeftModel.from_pretrained supports
+        memory_limits = {0: "0.4GIB", "cpu": "5GIB"}
+        # offloads around half of all transformer modules
+        device_map = infer_auto_device_map(model, max_memory=memory_limits)
+        self.assertTrue(0 in device_map.values())
+        self.assertTrue("cpu" in device_map.values())
+
+        config = LoraConfig(task_type="CAUSAL_LM", init_lora_weights=False, target_modules=["c_attn"])
+
+        model = get_peft_model(model, config)
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            model.save_pretrained(tmp_dir)
+            # load the model with device_map
+            model = AutoModelForCausalLM.from_pretrained(self.causal_lm_model_id, device_map=device_map).eval()
+            self.assertTrue(len({p.device for p in model.parameters()}) == 2)
+            model = PeftModel.from_pretrained(model, tmp_dir, max_memory=memory_limits)
+
+        input_tokens = tokenizer.encode("Four score and seven years ago", return_tensors="pt")
+        model.eval()
+
+        # test peft model adapter merge
+        pre_merge_olayer = model(input_tokens)[0]
+        model.merge_adapter()
+        post_merge_olayer = model(input_tokens)[0]
+        self.assertTrue(torch.allclose(post_merge_olayer, pre_merge_olayer))
+
+        # test peft model adapter unmerge
+        model.unmerge_adapter()
+        post_unmerge_olayer = model(input_tokens)[0]
+        self.assertTrue(torch.allclose(post_unmerge_olayer, pre_merge_olayer))
+
+        # test LoRA merge and unload
+        model = model.merge_and_unload()
+        post_unload_merge_olayer = model(input_tokens)[0]
+        self.assertTrue(torch.allclose(post_unload_merge_olayer, pre_merge_olayer))
+
+
 @require_torch_gpu
 class LoftQTests(unittest.TestCase):
     r"""