Llama bt update (#2700)

* adding Bt * adding optimum * fmt * revert transformers to 4.31.0 --------- Co-authored-by: lxning <[email protected]> Co-authored-by: lxning <[email protected]>
pytorch · Oct 11, 2023 · c69defd · c69defd
1 parent 9eb16cd
commit c69defd
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 8 deletions.
diff --git a/examples/large_models/Huggingface_accelerate/llama2/custom_handler.py b/examples/large_models/Huggingface_accelerate/llama2/custom_handler.py
@@ -3,9 +3,7 @@
 
 import torch
 import transformers
-from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
-from accelerate import init_empty_weights
-from accelerate import load_checkpoint_and_dispatch
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
 from ts.context import Context
 from ts.torch_handler.base_handler import BaseHandler
@@ -48,15 +46,24 @@ def initialize(self, ctx: Context):
             low_cpu_mem_usage=True,
             torch_dtype=torch.float16,
             load_in_8bit=True,
-            trust_remote_code=True)
-        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+            trust_remote_code=True,
+        )
+        if ctx.model_yaml_config["handler"]["fast_kernels"]:
+            from optimum.bettertransformer import BetterTransformer
+
+            try:
+                self.model = BetterTransformer.transform(self.model)
+            except RuntimeError as error:
+                logger.warning(
+                    "HuggingFace Optimum is not supporting this model,for the list of supported models, please refer to this doc,https://huggingface.co/docs/optimum/bettertransformer/overview"
+                )
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
         self.tokenizer.add_special_tokens(
             {
-
-            "pad_token": "<PAD>",
+                "pad_token": "<PAD>",
             }
         )
-        self.model.resize_token_embeddings(self.model.config.vocab_size + 1) 
+        self.model.resize_token_embeddings(self.model.config.vocab_size + 1)
 
         logger.info("Model %s loaded successfully", ctx.model_name)
         self.initialized = True

diff --git a/examples/large_models/Huggingface_accelerate/llama2/model-config.yaml b/examples/large_models/Huggingface_accelerate/llama2/model-config.yaml
@@ -11,3 +11,4 @@ handler:
     max_length: 50
     max_new_tokens: 50
     manual_seed: 40
+    fast_kernels: True
diff --git a/examples/large_models/Huggingface_accelerate/llama2/requirements.txt b/examples/large_models/Huggingface_accelerate/llama2/requirements.txt
@@ -3,3 +3,4 @@ accelerate
 bitsandbytes
 scipy
 mpi4py
+optimum