Merge pull request #3 from paperswithcode/feature/path_or_model

Marcin Kardas · web-flow · commit c818124435c1 · 2023-01-27T11:20:15.000Z
Feature/path or model
diff --git a/galai/__init__.py b/galai/__init__.py
@@ -4,7 +4,7 @@
 from galai.utils import ModelInfo
 import torch
 import warnings
-
+from pathlib import Path
 
 HF_MAPPING = {
     "mini": ("facebook/galactica-125m", torch.float32),
@@ -54,12 +54,19 @@ def load_model(
     Model - model object
     """
 
-    if name not in HF_MAPPING:
+    if name in HF_MAPPING:
+        hf_model, default_dtype = HF_MAPPING[name]
+        galai_model = True
+    elif Path(name).exists():
+        hf_model = name
+        default_dtype = torch.float32
+        galai_model = False
+    else:
         raise ValueError(
-            "Invalid model name. Must be one of 'mini', 'base', 'standard', 'large', 'huge'."
+            "Invalid model name. Must be one of 'mini', 'base', 'standard', 'large', 'huge', " +
+            "a path to a local checkpoint dir, or a model name available on HuggingFace hub."
         )
 
-    hf_model, default_dtype = HF_MAPPING[name]
     if dtype is None:
         dtype = default_dtype
 
@@ -99,7 +106,7 @@ def load_model(
                 UserWarning
             )
             num_gpus = available
-    if num_gpus > 1 and parallelize:
+    if num_gpus > 1 and parallelize and galai_model:
         mi = ModelInfo.by_name(name)
         if mi.num_heads % num_gpus != 0:
             raise ValueError(
diff --git a/galai/model.py b/galai/model.py
@@ -3,8 +3,7 @@
 
 import torch
 
-from tokenizers import Tokenizer
-from transformers import OPTForCausalLM, StoppingCriteriaList, StoppingCriteria
+from transformers import AutoTokenizer, OPTForCausalLM, StoppingCriteriaList, StoppingCriteria
 from parallelformers import parallelize
 import psutil
 
@@ -80,6 +79,7 @@ def __init__(
         self.is_loaded = False
         self.num_gpus = num_gpus
         self.tensor_parallel = tensor_parallel
+        self.max_input_length = 2020
         self._master_port = None
 
     def _load_checkpoint(self, checkpoint_path: str):
@@ -129,9 +129,15 @@ def _parallelize(self) -> None:
 
         self._master_port = 13000 + (id(self.model) % 32749)
 
+        custom_policies = None
+        if self.model.config.model_type == "opt" and not self.model.config.enable_bias:
+            from galai.parallel_policy import OPTDecoderLayerPolicyNoBias
+            custom_policies = [OPTDecoderLayerPolicyNoBias]
+
         parallelize(
             self.model, num_gpus=self.num_gpus, fp16=self.dtype == torch.float16,
             master_port=self._master_port,
+            custom_policies=custom_policies,
         )
 
     def _set_tokenizer(self, tokenizer_path: str):
@@ -143,9 +149,27 @@ def _set_tokenizer(self, tokenizer_path: str):
         tokenizer_path : str
             Path for the tokenizer (str)
         """
-        self.tokenizer = Tokenizer.from_pretrained(tokenizer_path)
-        self.tokenizer.enable_padding(direction="left", pad_id=1, pad_type_id=0, pad_token="[PAD]")
-        self.tokenizer.enable_truncation(max_length=2020, direction="left")
+        tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
+
+        # setup padding
+        tokenizer.pad_token_id = 1
+        tokenizer.pad_token = "<pad>"
+        tokenizer.padding_side = "left"
+
+        # setup truncation
+        tokenizer.truncation_side = "left"
+
+        # setup special tokens
+        tokenizer.bos_token_id = 0
+        tokenizer.bos_token = "<s>"
+
+        tokenizer.eos_token_id = 2
+        tokenizer.eos_token = "</s>"
+
+        tokenizer.unk_token = "<unk>"
+        tokenizer.unk_token_id = 3
+
+        self.tokenizer = tokenizer
 
     def _tokenize(self, input_text: List[str], new_doc: bool) -> torch.LongTensor:
         """
@@ -164,24 +188,27 @@ def _tokenize(self, input_text: List[str], new_doc: bool) -> torch.LongTensor:
             text = escape_custom_split_sequence(text)
             if not text:
                 warnings.warn(
-                    "Found an empty input text. Chainging to end-of-document token instead.",
+                    "Found an empty input text. Changing to end-of-document token instead.",
                     UserWarning
                 )
-                text = "</s>"
+                text = self.tokenizer.eos_token
             texts.append(text)
 
         if new_doc:
-            pad_id = self.tokenizer.padding["pad_id"]
-            pad_token = self.tokenizer.id_to_token(pad_id)
+            pad_token = self.tokenizer.pad_token
             texts = [pad_token + t for t in texts]
 
-        list_encoded = self.tokenizer.encode_batch(texts)
-        context_tokens = [encoded.ids for encoded in list_encoded]
+        encoded = self.tokenizer(
+            texts,
+            padding="longest",
+            max_length=self.max_input_length,
+            truncation=True
+        )
+        context_tokens = encoded["input_ids"]
         input_v = torch.LongTensor(context_tokens).to(self.model.device)
 
         if new_doc:
-            eos_id = self.tokenizer.token_to_id("</s>")
-            input_v[input_v[:, 0] == pad_id, 0] = eos_id
+            input_v[input_v[:, 0] == self.tokenizer.pad_token_id, 0] = self.tokenizer.eos_token_id
         return input_v
 
     @torch.inference_mode()
@@ -275,9 +302,12 @@ def generate(
         )
 
         # we keep special tokens such as [START_REF] or <work>
-        decoded = self.tokenizer.decode_batch(out['sequences'].tolist(), skip_special_tokens=False)
+        decoded = self.tokenizer.batch_decode(out['sequences'], skip_special_tokens=False)
         # so we manually remove </s> and <pad>
-        decoded = [text.replace("</s>", "").replace("<pad>", "") for text in decoded]
+        decoded = [
+            text.replace(self.tokenizer.eos_token, "").replace(self.tokenizer.pad_token, "")
+            for text in decoded
+        ]
 
         if num_return_sequences == 1:
             return decoded[0] if isinstance(input_text, str) else decoded
@@ -363,7 +393,7 @@ def generate_reference(
         prompt_length = input_v.shape[1]
         finished_reference_criteria = FinishedReferenceCriteria(
             prompt_length=prompt_length,
-            end_ref_id=self.tokenizer.token_to_id("[END_REF]"),
+            end_ref_id=self.tokenizer.convert_tokens_to_ids("[END_REF]"),
         )
 
         if max_new_tokens is None and max_length is None:
@@ -396,8 +426,8 @@ def generate_reference(
                 stopping_criteria=stopping_criteria,
             )
         # cut-off the prompts
-        generated_tokens = out["sequences"][:, prompt_length:].tolist()
-        decoded = self.tokenizer.decode_batch(generated_tokens, skip_special_tokens=False)
+        generated_tokens = out["sequences"][:, prompt_length:]
+        decoded = self.tokenizer.batch_decode(generated_tokens, skip_special_tokens=False)
         references = []
         unfinished_generation = False
         for text in decoded:
diff --git a/galai/parallel_policy.py b/galai/parallel_policy.py
@@ -0,0 +1,60 @@
+from parallelformers.policies.base import Layer, Policy
+from parallelformers.utils.dist_utils import AllReduceLinear
+
+from transformers.models.opt.modeling_opt import OPTDecoderLayer
+
+
+__all__ = ["OPTDecoderLayerPolicyNoBias"]
+
+
+class OPTDecoderLayerPolicyNoBias(Policy):
+    @staticmethod
+    def replace_arguments(config, world_size):
+        return {
+            "self_attn.embed_dim": config.hidden_size // world_size,
+            "self_attn.num_heads": config.num_attention_heads // world_size,
+        }
+
+    @staticmethod
+    def attn_qkv():
+        return [
+            Layer(
+                weight="self_attn.q_proj.weight",
+            ),
+            Layer(
+                weight="self_attn.k_proj.weight",
+            ),
+            Layer(
+                weight="self_attn.v_proj.weight",
+            ),
+        ]
+
+    @staticmethod
+    def attn_out():
+        return [
+            Layer(
+                weight="self_attn.out_proj.weight",
+                replace=AllReduceLinear,
+            ),
+        ]
+
+    @staticmethod
+    def mlp_in():
+        return [
+            Layer(
+                weight="fc1.weight",
+            ),
+        ]
+
+    @staticmethod
+    def mlp_out():
+        return [
+            Layer(
+                weight="fc2.weight",
+                replace=AllReduceLinear,
+            ),
+        ]
+
+    @staticmethod
+    def original_layer_class():
+        return OPTDecoderLayer
