[Model] Udpate default prefill chunk size and max batch size (mlc-ai#…

…2917) This PR updates the default prefill chunk size from 2048 to 8192, and the default max batch size from 80 to 128.
cjmcv · Sep 18, 2024 · 66fd62d · 66fd62d
1 parent 1eabc65
commit 66fd62d
Show file tree

Hide file tree

Showing 21 changed files with 79 additions and 79 deletions.
diff --git a/python/mlc_llm/cli/gen_config.py b/python/mlc_llm/cli/gen_config.py
@@ -86,7 +86,7 @@ def _parse_output(path: Union[str, Path]) -> Path:
     parser.add_argument(
         "--max-batch-size",
         type=int,
-        default=80,
+        default=128,
         help=HELP["max_batch_size"] + ' (default: "%(default)s")',
     )
     parser.add_argument(

diff --git a/python/mlc_llm/model/baichuan/baichuan_model.py b/python/mlc_llm/model/baichuan/baichuan_model.py
@@ -68,17 +68,17 @@ def __post_init__(self):
             logger.info(
                 "%s defaults to %d",
                 bold("prefill_chunk_size"),
-                min(self.context_window_size, 2048),
+                min(self.context_window_size, 8192),
             )
-            self.prefill_chunk_size = min(self.context_window_size, 2048)
+            self.prefill_chunk_size = min(self.context_window_size, 8192)
         elif self.prefill_chunk_size > self.context_window_size:
             logger.info(
                 "Overriding %s from %d to %d",
                 bold("prefill_chunk_size"),
                 self.prefill_chunk_size,
-                min(self.context_window_size, 2048),
+                min(self.context_window_size, 8192),
             )
-            self.prefill_chunk_size = min(self.context_window_size, 2048)
+            self.prefill_chunk_size = min(self.context_window_size, 8192)
 
 
 # pylint: disable=invalid-name,missing-docstring

diff --git a/python/mlc_llm/model/chatglm3/chatglm3_model.py b/python/mlc_llm/model/chatglm3/chatglm3_model.py
@@ -74,17 +74,17 @@ def __post_init__(self):
             logger.info(
                 "%s defaults to %d",
                 bold("prefill_chunk_size"),
-                min(self.context_window_size, 2048),
+                min(self.context_window_size, 8192),
             )
-            self.prefill_chunk_size = min(self.context_window_size, 2048)
+            self.prefill_chunk_size = min(self.context_window_size, 8192)
         elif self.prefill_chunk_size > self.context_window_size:
             logger.info(
                 "Overriding %s from %d to %d",
                 bold("prefill_chunk_size"),
                 self.prefill_chunk_size,
-                min(self.context_window_size, 2048),
+                min(self.context_window_size, 8192),
             )
-            self.prefill_chunk_size = min(self.context_window_size, 2048)
+            self.prefill_chunk_size = min(self.context_window_size, 8192)
 
 
 # pylint: disable=invalid-name,missing-docstring

diff --git a/python/mlc_llm/model/cohere/cohere_model.py b/python/mlc_llm/model/cohere/cohere_model.py
@@ -63,17 +63,17 @@ def __post_init__(self):
                 logger.info(
                     "%s defaults to %d",
                     bold("prefill_chunk_size"),
-                    min(self.context_window_size, 2048),
+                    min(self.context_window_size, 8192),
                 )
-                self.prefill_chunk_size = min(self.context_window_size, 2048)
+                self.prefill_chunk_size = min(self.context_window_size, 8192)
             elif self.prefill_chunk_size > self.context_window_size:
                 logger.info(
                     "Overriding %s from %d to %d",
                     bold("prefill_chunk_size"),
                     self.prefill_chunk_size,
-                    min(self.context_window_size, 2048),
+                    min(self.context_window_size, 8192),
                 )
-                self.prefill_chunk_size = min(self.context_window_size, 2048)
+                self.prefill_chunk_size = min(self.context_window_size, 8192)
 
             if self.num_key_value_heads == 0 or self.num_key_value_heads is None:
                 self.num_key_value_heads = self.num_attention_heads

diff --git a/python/mlc_llm/model/gemma/gemma_model.py b/python/mlc_llm/model/gemma/gemma_model.py
@@ -72,17 +72,17 @@ def __post_init__(self):
             logger.info(
                 "%s defaults to %d",
                 bold("prefill_chunk_size"),
-                min(self.context_window_size, 2048),
+                min(self.context_window_size, 8192),
             )
-            self.prefill_chunk_size = min(self.context_window_size, 2048)
+            self.prefill_chunk_size = min(self.context_window_size, 8192)
         elif self.prefill_chunk_size > self.context_window_size:
             logger.info(
                 "Overriding %s from %d to %d",
                 bold("prefill_chunk_size"),
                 self.prefill_chunk_size,
-                min(self.context_window_size, 2048),
+                min(self.context_window_size, 8192),
             )
-            self.prefill_chunk_size = min(self.context_window_size, 2048)
+            self.prefill_chunk_size = min(self.context_window_size, 8192)
 
 
 # pylint: disable=invalid-name,missing-docstring

diff --git a/python/mlc_llm/model/gpt2/gpt2_model.py b/python/mlc_llm/model/gpt2/gpt2_model.py
@@ -65,17 +65,17 @@ def __post_init__(self):
             logger.info(
                 "%s defaults to %d",
                 bold("prefill_chunk_size"),
-                min(self.context_window_size, 2048),
+                min(self.context_window_size, 8192),
             )
-            self.prefill_chunk_size = min(self.context_window_size, 2048)
+            self.prefill_chunk_size = min(self.context_window_size, 8192)
         elif self.prefill_chunk_size > self.context_window_size:
             logger.info(
                 "Overriding %s from %d to %d",
                 bold("prefill_chunk_size"),
                 self.prefill_chunk_size,
-                min(self.context_window_size, 2048),
+                min(self.context_window_size, 8192),
             )
-            self.prefill_chunk_size = min(self.context_window_size, 2048)
+            self.prefill_chunk_size = min(self.context_window_size, 8192)
 
 
 # pylint: disable=invalid-name,missing-docstring,too-many-locals

diff --git a/python/mlc_llm/model/gpt_bigcode/gpt_bigcode_model.py b/python/mlc_llm/model/gpt_bigcode/gpt_bigcode_model.py
@@ -57,17 +57,17 @@ def __post_init__(self):
             logger.info(
                 "%s defaults to %d",
                 bold("prefill_chunk_size"),
-                min(self.context_window_size, 2048),
+                min(self.context_window_size, 8192),
             )
-            self.prefill_chunk_size = min(self.context_window_size, 2048)
+            self.prefill_chunk_size = min(self.context_window_size, 8192)
         elif self.prefill_chunk_size > self.context_window_size:
             logger.info(
                 "Overriding %s from %d to %d",
                 bold("prefill_chunk_size"),
                 self.prefill_chunk_size,
-                min(self.context_window_size, 2048),
+                min(self.context_window_size, 8192),
             )
-            self.prefill_chunk_size = min(self.context_window_size, 2048)
+            self.prefill_chunk_size = min(self.context_window_size, 8192)
 
 
 # pylint: disable=invalid-name,missing-docstring

diff --git a/python/mlc_llm/model/gpt_neox/gpt_neox_model.py b/python/mlc_llm/model/gpt_neox/gpt_neox_model.py
@@ -72,17 +72,17 @@ def __post_init__(self):
             logger.info(
                 "%s defaults to %d",
                 bold("prefill_chunk_size"),
-                min(self.context_window_size, 2048),
+                min(self.context_window_size, 8192),
             )
-            self.prefill_chunk_size = min(self.context_window_size, 2048)
+            self.prefill_chunk_size = min(self.context_window_size, 8192)
         elif self.prefill_chunk_size > self.context_window_size:
             logger.info(
                 "Overriding %s from %d to %d",
                 bold("prefill_chunk_size"),
                 self.prefill_chunk_size,
-                min(self.context_window_size, 2048),
+                min(self.context_window_size, 8192),
             )
-            self.prefill_chunk_size = min(self.context_window_size, 2048)
+            self.prefill_chunk_size = min(self.context_window_size, 8192)
 
 
 # pylint: disable=invalid-name,missing-docstring

diff --git a/python/mlc_llm/model/internlm/internlm_model.py b/python/mlc_llm/model/internlm/internlm_model.py
@@ -67,17 +67,17 @@ def __post_init__(self):
             logger.info(
                 "%s defaults to %d",
                 bold("prefill_chunk_size"),
-                min(self.context_window_size, 2048),
+                min(self.context_window_size, 8192),
             )
-            self.prefill_chunk_size = min(self.context_window_size, 2048)
+            self.prefill_chunk_size = min(self.context_window_size, 8192)
         elif self.prefill_chunk_size > self.context_window_size:
             logger.info(
                 "Overriding %s from %d to %d",
                 bold("prefill_chunk_size"),
                 self.prefill_chunk_size,
-                min(self.context_window_size, 2048),
+                min(self.context_window_size, 8192),
             )
-            self.prefill_chunk_size = min(self.context_window_size, 2048)
+            self.prefill_chunk_size = min(self.context_window_size, 8192)
 
 
 # pylint: disable=invalid-name,missing-docstring

diff --git a/python/mlc_llm/model/internlm2/internlm2_model.py b/python/mlc_llm/model/internlm2/internlm2_model.py
@@ -69,17 +69,17 @@ def __post_init__(self):
             logger.info(
                 "%s defaults to %d",
                 bold("prefill_chunk_size"),
-                min(self.context_window_size, 2048),
+                min(self.context_window_size, 8192),
             )
-            self.prefill_chunk_size = min(self.context_window_size, 2048)
+            self.prefill_chunk_size = min(self.context_window_size, 8192)
         elif self.prefill_chunk_size > self.context_window_size:
             logger.info(
                 "Overriding %s from %d to %d",
                 bold("prefill_chunk_size"),
                 self.prefill_chunk_size,
-                min(self.context_window_size, 2048),
+                min(self.context_window_size, 8192),
             )
-            self.prefill_chunk_size = min(self.context_window_size, 2048)
+            self.prefill_chunk_size = min(self.context_window_size, 8192)
 
 
 # pylint: disable=invalid-name,missing-docstring

diff --git a/python/mlc_llm/model/llama/llama_model.py b/python/mlc_llm/model/llama/llama_model.py
@@ -89,17 +89,17 @@ def __post_init__(self):  # pylint: disable=too-many-branches
             logger.info(
                 "%s defaults to %d",
                 bold("prefill_chunk_size"),
-                min(self.context_window_size, 2048),
+                min(self.context_window_size, 8192),
             )
-            self.prefill_chunk_size = min(self.context_window_size, 2048)
+            self.prefill_chunk_size = min(self.context_window_size, 8192)
         elif self.prefill_chunk_size > self.context_window_size:
             logger.info(
                 "Overriding %s from %d to %d",
                 bold("prefill_chunk_size"),
                 self.prefill_chunk_size,
-                min(self.context_window_size, 2048),
+                min(self.context_window_size, 8192),
             )
-            self.prefill_chunk_size = min(self.context_window_size, 2048)
+            self.prefill_chunk_size = min(self.context_window_size, 8192)
 
 
 # pylint: disable=invalid-name,missing-docstring

diff --git a/python/mlc_llm/model/minicpm/minicpm_model.py b/python/mlc_llm/model/minicpm/minicpm_model.py
@@ -77,17 +77,17 @@ def __post_init__(self):
             logger.info(
                 "%s defaults to %d",
                 bold("prefill_chunk_size"),
-                min(self.context_window_size, 2048),
+                min(self.context_window_size, 8192),
             )
-            self.prefill_chunk_size = min(self.context_window_size, 2048)
+            self.prefill_chunk_size = min(self.context_window_size, 8192)
         elif self.prefill_chunk_size > self.context_window_size:
             logger.info(
                 "Overriding %s from %d to %d",
                 bold("prefill_chunk_size"),
                 self.prefill_chunk_size,
-                min(self.context_window_size, 2048),
+                min(self.context_window_size, 8192),
             )
-            self.prefill_chunk_size = min(self.context_window_size, 2048)
+            self.prefill_chunk_size = min(self.context_window_size, 8192)
 
 
 # pylint: disable=invalid-name,missing-docstring

diff --git a/python/mlc_llm/model/mistral/mistral_model.py b/python/mlc_llm/model/mistral/mistral_model.py
@@ -87,9 +87,9 @@ def __post_init__(self):  # pylint: disable=too-many-branches
             logger.info(
                 "%s defaults to %d",
                 bold("prefill_chunk_size"),
-                min(*prefill_chunk_size_candidates, 2048),
+                min(*prefill_chunk_size_candidates, 8192),
             )
-            self.prefill_chunk_size = min(*prefill_chunk_size_candidates, 2048)
+            self.prefill_chunk_size = min(*prefill_chunk_size_candidates, 8192)
 
 
 # pylint: disable=invalid-name,missing-docstring

diff --git a/python/mlc_llm/model/orion/orion_model.py b/python/mlc_llm/model/orion/orion_model.py
@@ -72,17 +72,17 @@ def __post_init__(self):
             logger.info(
                 "%s defaults to %d",
                 bold("prefill_chunk_size"),
-                min(self.context_window_size, 2048),
+                min(self.context_window_size, 8192),
             )
-            self.prefill_chunk_size = min(self.context_window_size, 2048)
+            self.prefill_chunk_size = min(self.context_window_size, 8192)
         elif self.prefill_chunk_size > self.context_window_size:
             logger.info(
                 "Overriding %s from %d to %d",
                 bold("prefill_chunk_size"),
                 self.prefill_chunk_size,
-                min(self.context_window_size, 2048),
+                min(self.context_window_size, 8192),
             )
-            self.prefill_chunk_size = min(self.context_window_size, 2048)
+            self.prefill_chunk_size = min(self.context_window_size, 8192)
 
 
 # pylint: disable=invalid-name,missing-docstring

diff --git a/python/mlc_llm/model/phi/phi_model.py b/python/mlc_llm/model/phi/phi_model.py
@@ -67,17 +67,17 @@ def __post_init__(self):
             logger.info(
                 "%s defaults to %d",
                 bold("prefill_chunk_size"),
-                min(self.context_window_size, 2048),
+                min(self.context_window_size, 8192),
             )
-            self.prefill_chunk_size = min(self.context_window_size, 2048)
+            self.prefill_chunk_size = min(self.context_window_size, 8192)
         elif self.prefill_chunk_size > self.context_window_size:
             logger.info(
                 "Overriding %s from %d to %d",
                 bold("prefill_chunk_size"),
                 self.prefill_chunk_size,
-                min(self.context_window_size, 2048),
+                min(self.context_window_size, 8192),
             )
-            self.prefill_chunk_size = min(self.context_window_size, 2048)
+            self.prefill_chunk_size = min(self.context_window_size, 8192)
         if self.num_key_value_heads == 0 or self.num_key_value_heads is None:
             self.num_key_value_heads = self.num_attention_heads
         if self.intermediate_size == 0 or self.intermediate_size is None:

diff --git a/python/mlc_llm/model/phi3/phi3_model.py b/python/mlc_llm/model/phi3/phi3_model.py
@@ -72,17 +72,17 @@ def __post_init__(self):
             logger.info(
                 "%s defaults to %d",
                 bold("prefill_chunk_size"),
-                min(self.context_window_size, 2048),
+                min(self.context_window_size, 8192),
             )
-            self.prefill_chunk_size = min(self.context_window_size, 2048)
+            self.prefill_chunk_size = min(self.context_window_size, 8192)
         elif self.prefill_chunk_size > self.context_window_size:
             logger.info(
                 "Overriding %s from %d to %d",
                 bold("prefill_chunk_size"),
                 self.prefill_chunk_size,
-                min(self.context_window_size, 2048),
+                min(self.context_window_size, 8192),
             )
-            self.prefill_chunk_size = min(self.context_window_size, 2048)
+            self.prefill_chunk_size = min(self.context_window_size, 8192)
 
         if self.num_key_value_heads == 0 or self.num_key_value_heads is None:
             self.num_key_value_heads = self.num_attention_heads

diff --git a/python/mlc_llm/model/phi3v/phi3v_model.py b/python/mlc_llm/model/phi3v/phi3v_model.py
@@ -101,17 +101,17 @@ def __post_init__(self):
             logger.info(
                 "%s defaults to %d",
                 bold("prefill_chunk_size"),
-                min(self.context_window_size, 2048),
+                min(self.context_window_size, 8192),
             )
-            self.prefill_chunk_size = min(self.context_window_size, 2048)
+            self.prefill_chunk_size = min(self.context_window_size, 8192)
         elif self.prefill_chunk_size > self.context_window_size:
             logger.info(
                 "Overriding %s from %d to %d",
                 bold("prefill_chunk_size"),
                 self.prefill_chunk_size,
-                min(self.context_window_size, 2048),
+                min(self.context_window_size, 8192),
             )
-            self.prefill_chunk_size = min(self.context_window_size, 2048)
+            self.prefill_chunk_size = min(self.context_window_size, 8192)
 
         if self.num_key_value_heads == 0 or self.num_key_value_heads is None:
             self.num_key_value_heads = self.num_attention_heads

diff --git a/python/mlc_llm/model/qwen/qwen_model.py b/python/mlc_llm/model/qwen/qwen_model.py
@@ -65,17 +65,17 @@ def __post_init__(self):
             logger.info(
                 "%s defaults to %d",
                 bold("prefill_chunk_size"),
-                min(self.context_window_size, 2048),
+                min(self.context_window_size, 8192),
             )
-            self.prefill_chunk_size = min(self.context_window_size, 2048)
+            self.prefill_chunk_size = min(self.context_window_size, 8192)
         elif self.prefill_chunk_size > self.context_window_size:
             logger.info(
                 "Overriding %s from %d to %d",
                 bold("prefill_chunk_size"),
                 self.prefill_chunk_size,
-                min(self.context_window_size, 2048),
+                min(self.context_window_size, 8192),
             )
-            self.prefill_chunk_size = min(self.context_window_size, 2048)
+            self.prefill_chunk_size = min(self.context_window_size, 8192)
 
 
 # pylint: disable=invalid-name,missing-docstring

diff --git a/python/mlc_llm/model/qwen2/qwen2_model.py b/python/mlc_llm/model/qwen2/qwen2_model.py
@@ -67,17 +67,17 @@ def __post_init__(self):
             logger.info(
                 "%s defaults to %d",
                 bold("prefill_chunk_size"),
-                min(self.context_window_size, 2048),
+                min(self.context_window_size, 8192),
             )
-            self.prefill_chunk_size = min(self.context_window_size, 2048)
+            self.prefill_chunk_size = min(self.context_window_size, 8192)
         elif self.prefill_chunk_size > self.context_window_size:
             logger.info(
                 "Overriding %s from %d to %d",
                 bold("prefill_chunk_size"),
                 self.prefill_chunk_size,
-                min(self.context_window_size, 2048),
+                min(self.context_window_size, 8192),
             )
-            self.prefill_chunk_size = min(self.context_window_size, 2048)
+            self.prefill_chunk_size = min(self.context_window_size, 8192)
 
 
 # pylint: disable=invalid-name,missing-docstring,too-many-locals