wassname/peft
1
0
Fork 0
mirror of https://github.com/wassname/peft.git synced 2026-09-09 11:28:32 +08:00
Code Issues Packages Projects Releases Wiki Activity

update peft_model_load_and_dispatch

Browse Source
This commit is contained in:
Sourab Mangrulkar
2023-02-08 09:21:49 +05:30
parent a7dd034710
commit 64d2d19598
1 changed files with 3 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/peft/utils/save_and_load.py
+3 -4
View File
@@ -79,7 +79,7 @@ def set_peft_model_state_dict(model, peft_model_state_dict):
return model
def peft_model_load_and_dispatch(model, peft_model_state_dict, peft_config, max_memory=None):
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.
@@ -94,12 +94,11 @@ def peft_model_load_and_dispatch(model, peft_model_state_dict, peft_config, max_
from accelerate import dispatch_model, infer_auto_device_map
from accelerate.hooks import AlignDevicesHook, add_hook_to_module, remove_hook_from_submodules
from ..mapping import get_peft_model
from ..peft_model import PeftModel
remove_hook_from_submodules(model)
model = get_peft_model(model, peft_config)
model = PeftModel.from_pretrained(peft_model_id, peft_model_id)
model.print_trainable_parameters()
set_peft_model_state_dict(model, peft_model_state_dict)
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)