Merge pull request #418 from stevelaskaridis/dev/hf-device-map

Propagate device_map to hf model

dsp/modules/hf.py

@@ -37,7 +37,7 @@ def __init__(self, model: str, checkpoint: Optional[str] = None, is_client: bool
             model (str): HF model identifier to load and use
             checkpoint (str, optional): load specific checkpoints of the model. Defaults to None.
             is_client (bool, optional): whether to access models via client. Defaults to False.
-            hf_device_map (str, optional): HF config strategy to load the model. 
+            hf_device_map (str, optional): HF config strategy to load the model.
                 Recommeded to use "auto", which will help loading large models using accelerate. Defaults to "auto".
         """
 
@@ -72,14 +72,20 @@ def __init__(self, model: str, checkpoint: Optional[str] = None, is_client: bool
                     #     self.model = AutoModelClass.from_pretrained(peft_config.base_model_name_or_path, return_dict=True, load_in_8bit=True, device_map=hf_device_map)
                     #     self.model = PeftModel.from_pretrained(self.model, checkpoint)
                     # else:
-                    self.model = AutoModelClass.from_pretrained(checkpoint).to(self.device)
+                    if self.device_map:
+                        self.model = AutoModelClass.from_pretrained(checkpoint, device_map=self.device_map)
+                    else:
+                        self.model = AutoModelClass.from_pretrained(checkpoint).to(self.device)
                 else:
-                    self.model = AutoModelClass.from_pretrained(model).to(self.device)
+                    if self.device_map:
+                        self.model = AutoModelClass.from_pretrained(model, device_map=self.device_map)
+                    else:
+                        self.model = AutoModelClass.from_pretrained(model).to(self.device)
                 self.drop_prompt_from_output = False
             except ValueError:
                 self.model = AutoModelForCausalLM.from_pretrained(
                     model if checkpoint is None else checkpoint,
-                    device_map=hf_device_map
+                    device_map=self.device_map
                 )
                 self.drop_prompt_from_output = True
                 self.tokenizer = AutoTokenizer.from_pretrained(model)