From fc7f166f00b7fa6f1d704b7bd21ff77557d245ed Mon Sep 17 00:00:00 2001 From: Ben Steel Date: Tue, 21 Nov 2023 18:03:36 -0500 Subject: [PATCH 1/3] Add the ability to set model kwargs for HF local models --- dsp/modules/hf.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/dsp/modules/hf.py b/dsp/modules/hf.py index 79cdffeb1c..78ace1fbe0 100644 --- a/dsp/modules/hf.py +++ b/dsp/modules/hf.py @@ -30,7 +30,8 @@ def openai_to_hf(**kwargs): class HFModel(LM): def __init__(self, model: str, checkpoint: Optional[str] = None, is_client: bool = False, - hf_device_map: Literal["auto", "balanced", "balanced_low_0", "sequential"] = "auto"): + hf_device_map: Literal["auto", "balanced", "balanced_low_0", "sequential"] = "auto", + model_kwargs: Optional[dict] = None): """wrapper for Hugging Face models Args: @@ -39,6 +40,7 @@ def __init__(self, model: str, checkpoint: Optional[str] = None, is_client: bool 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. Recommeded to use "auto", which will help loading large models using accelerate. Defaults to "auto". + model_kwargs (dict, optional): additional kwargs to pass to the model. Defaults to None. """ try: from transformers import AutoModelForSeq2SeqLM, AutoModelForCausalLM, AutoTokenizer, AutoConfig @@ -71,14 +73,15 @@ 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("cuda") + self.model = AutoModelClass.from_pretrained(checkpoint, **model_kwargs).to("cuda") else: - self.model = AutoModelClass.from_pretrained(model).to("cuda") + self.model = AutoModelClass.from_pretrained(model, **model_kwargs).to("cuda") 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=hf_device_map, + **model_kwargs, ) self.drop_prompt_from_output = True self.tokenizer = AutoTokenizer.from_pretrained(model) From 430e92bb48aff7f5d58dcb7a0e737d25e21bd093 Mon Sep 17 00:00:00 2001 From: Ben Steel Date: Tue, 5 Dec 2023 21:53:16 -0500 Subject: [PATCH 2/3] Fixing default value of model kwargs --- dsp/modules/hf.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dsp/modules/hf.py b/dsp/modules/hf.py index 78ace1fbe0..13c7ad6d45 100644 --- a/dsp/modules/hf.py +++ b/dsp/modules/hf.py @@ -31,7 +31,7 @@ def openai_to_hf(**kwargs): class HFModel(LM): def __init__(self, model: str, checkpoint: Optional[str] = None, is_client: bool = False, hf_device_map: Literal["auto", "balanced", "balanced_low_0", "sequential"] = "auto", - model_kwargs: Optional[dict] = None): + model_kwargs: Optional[dict] = {}): """wrapper for Hugging Face models Args: From f9a036271670f1a6833003d493dc6cf10bf565a4 Mon Sep 17 00:00:00 2001 From: Ben Steel Date: Mon, 15 Apr 2024 19:22:35 +0200 Subject: [PATCH 3/3] Correcting device_map conflict in local model extra kwargs --- dsp/modules/hf.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/dsp/modules/hf.py b/dsp/modules/hf.py index 3eae1bc004..cefcee9290 100644 --- a/dsp/modules/hf.py +++ b/dsp/modules/hf.py @@ -61,6 +61,10 @@ def __init__( hf_autoconfig_kwargs = dict(token=token or os.environ.get("HF_TOKEN")) hf_autotokenizer_kwargs = hf_autoconfig_kwargs.copy() hf_automodel_kwargs = hf_autoconfig_kwargs.copy() + + # silently remove device_map from model_kwargs if it is present, as the option is provided in the constructor + if "device_map" in model_kwargs: + model_kwargs.pop("device_map") hf_automodel_kwargs.update(model_kwargs) if not self.is_client: try: