Merge pull request #187 from OpenAccess-AI-Collective/strip-peft-device-map

src/axolotl/utils/models.py

@@ -364,7 +364,6 @@ def load_llama_adapter(model, cfg):
         model = PeftModel.from_pretrained(
             model,
             cfg.lora_model_dir,
-            device_map=cfg.device_map,
             torch_dtype=torch.float16,
         )
     else:
@@ -427,8 +426,6 @@ def load_lora(model, cfg):
             model,
             cfg.lora_model_dir,
             is_trainable=not cfg.inference,
-            device_map=cfg.device_map,
-            # torch_dtype=torch.float16,
         )
     else:
         model = get_peft_model(model, lora_config)