adjust the flashinfer llama model to accomodate baichuan

mlsys-io · Jul 7, 2024 · 5004444 · 5004444
1 parent d8dcdeb
commit 5004444
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 19 deletions.
diff --git a/server/examples/test_local_api.py b/server/examples/test_local_api.py
@@ -251,7 +251,9 @@ def make_input(lora_id, lora_or_base, id=0, promptOverride=None):
             promptOverride="What are the differences between Manhattan and Brooklyn",
         ),
     ]
-    service = FlashinferLlama(model_id="baichuan-inc/Baichuan2-7B-Chat")
+    service = FlashinferLlama(
+        model_id="baichuan-inc/Baichuan2-7B-Chat", trust_remote_code=True
+    )
 
 print(service.get_lora_adapters())
 tokenizer = service.tokenizer

diff --git a/server/text_generation_server/layers/flashinfer_attention.py b/server/text_generation_server/layers/flashinfer_attention.py
@@ -19,7 +19,7 @@ class AttentionRotaryParams:
     causal: bool = True
     pos_encoding_mode: POS_ENCODING_MODE = POS_ENCODING_MODE.ROPE_LLAMA
     rope_scale: float = 1.0
-    rope_theta: float = 1.0e-4
+    rope_theta: float = 1.0e4
 
 
 def find_padded_head_dim(head_dim):

diff --git a/server/text_generation_server/models/flashinfer_llama.py b/server/text_generation_server/models/flashinfer_llama.py
@@ -33,23 +33,14 @@ def __init__(
         else:
             raise NotImplementedError("Flashinfer Llama is only available on Cuda")
 
-        tokenizer = AutoTokenizer.from_pretrained(model_id)
-        try:
-            tokenizer = LlamaTokenizer.from_pretrained(
-                model_id,
-                revision=revision,
-                padding_side="left",
-                truncation_side="left",
-                trust_remote_code=trust_remote_code,
-            )
-        except Exception:
-            tokenizer = AutoTokenizer.from_pretrained(
-                model_id,
-                revision=revision,
-                padding_side="left",
-                truncation_side="left",
-                trust_remote_code=trust_remote_code,
-            )
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_id,
+            revision=revision,
+            padding_side="left",
+            truncation_side="left",
+            trust_remote_code=trust_remote_code,
+        )
+
         try:
             generation_config = GenerationConfig.from_pretrained(
                 model_id, revision=revision, trust_remote_code=trust_remote_code
@@ -65,6 +56,12 @@ def __init__(
         )
         config.quantize = quantize
         config.speculator = None
+        if not hasattr(config, "num_key_value_heads"):
+            config.num_key_value_heads = config.num_attention_heads
+
+        if not hasattr(config, "rope_theta"):
+            config.rope_theta = 1.0e4
+
         torch.distributed.barrier(group=self.process_group)
 
         filenames = weight_files(model_id, revision=revision, extension=".safetensors")