From 9bf6ecf9e2de9b72c3fa62e0e6f5b5e9041825b1 Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Thu, 23 Mar 2023 16:49:41 -0300 Subject: [PATCH] Fix LoRA device map (attempt) --- modules/LoRA.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/modules/LoRA.py b/modules/LoRA.py index 0a2aaa7d..5f77e340 100644 --- a/modules/LoRA.py +++ b/modules/LoRA.py @@ -16,12 +16,15 @@ def add_lora_to_model(lora_name): print(f"Adding the LoRA {lora_name} to the model...") params = {} - if shared.args.load_in_8bit: - params['device_map'] = {'': 0} - elif not shared.args.cpu: - params['device_map'] = 'auto' + if not shared.args.cpu: params['dtype'] = shared.model.dtype + if hasattr(shared.model, "hf_device_map"): + params['device_map'] = {"base_model.model."+k: v for k, v in shared.model.hf_device_map.items()} + elif shared.args.load_in_8bit: + params['device_map'] = {'': 0} shared.model = PeftModel.from_pretrained(shared.model, Path(f"loras/{lora_name}"), **params) if not shared.args.load_in_8bit and not shared.args.cpu: shared.model.half() + if not hasattr(shared.model, "hf_device_map"): + shared.model.cuda()