From 44d030c23630ecab67b1a8f744d4cb737d0a7665 Mon Sep 17 00:00:00 2001 From: Khai Loong Aw Date: Tue, 11 Jul 2023 09:22:50 +0200 Subject: [PATCH] Add code for multi-gpu HuggingfaceSubject --- brainscore_language/model_helpers/huggingface.py | 3 +-- pyproject.toml | 1 + 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/brainscore_language/model_helpers/huggingface.py b/brainscore_language/model_helpers/huggingface.py index 91338560..cb3af02d 100644 --- a/brainscore_language/model_helpers/huggingface.py +++ b/brainscore_language/model_helpers/huggingface.py @@ -42,9 +42,8 @@ def __init__( self._logger = logging.getLogger(fullname(self)) self.model_id = model_id self.region_layer_mapping = region_layer_mapping - self.basemodel = (model if model is not None else AutoModelForCausalLM.from_pretrained(self.model_id)) + self.basemodel = (model if model is not None else AutoModelForCausalLM.from_pretrained(self.model_id, device_map="auto")) self.device = 'cuda' if torch.cuda.is_available() else 'cpu' - self.basemodel.to(self.device) self.tokenizer = tokenizer if tokenizer is not None else AutoTokenizer.from_pretrained(self.model_id, truncation_side='left') self.current_tokens = None # keep track of current tokens diff --git a/pyproject.toml b/pyproject.toml index 78178129..4b6da964 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,6 +21,7 @@ dependencies = [ "transformers>=4.11.3", "gensim", "joblib", + "accelerate", ] [project.optional-dependencies]