From 1929a84e1e1f9f1dd2ea22225f92191a206b8d49 Mon Sep 17 00:00:00 2001
From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com>
Date: Wed, 8 Feb 2023 12:29:03 +0530
Subject: [PATCH] remove `peft_model_load_and_dispatch` as it is part of
 `PeftModel.from_pretrained`

---
 src/peft/utils/save_and_load.py | 31 -------------------------------
 1 file changed, 31 deletions(-)

diff --git a/src/peft/utils/save_and_load.py b/src/peft/utils/save_and_load.py
index b2f789e..c6596c7 100644
--- a/src/peft/utils/save_and_load.py
+++ b/src/peft/utils/save_and_load.py
@@ -77,34 +77,3 @@ def set_peft_model_state_dict(model, peft_model_state_dict):
             {"weight": peft_model_state_dict["prompt_embeddings"]}, strict=True
         )
     return model
-
-
-def peft_model_load_and_dispatch(model, peft_model_id, max_memory=None):
-    """
-    Load the Peft model state dict and dispatch the model to the correct device.
-
-    Args:
-        model ([`PeftModel`]): The Pre-trained base model which has already been sharded and dispatched
-        using `accelerate` functionalities.
-        peft_model_state_dict (`dict`): The state dict of the Peft model.
-        max_memory (`Dict`, *optional*):
-            A dictionary device identifier to maximum memory. Will default to the maximum memory available for each GPU
-            and the available CPU RAM if unset.
-    """
-    from accelerate import dispatch_model, infer_auto_device_map
-    from accelerate.hooks import AlignDevicesHook, add_hook_to_module, remove_hook_from_submodules
-
-    from ..peft_model import PeftModel
-
-    remove_hook_from_submodules(model)
-    model = PeftModel.from_pretrained(model, peft_model_id)
-    model.print_trainable_parameters()
-    device_map = infer_auto_device_map(model, max_memory=max_memory, no_split_module_classes=model._no_split_modules)
-    model = dispatch_model(model, device_map=device_map)
-    hook = AlignDevicesHook(io_same_device=True)
-    if model.peft_config.peft_type == PeftType.LORA:
-        add_hook_to_module(model.base_model.model, hook)
-    else:
-        remove_hook_from_submodules(model.prompt_encoder)
-        add_hook_to_module(model.base_model, hook)
-    return model