Add --num_experts_per_token parameter (ExLlamav2) (#4955)

oobabooga · Dec 17, 2023 · f1f2c4c · f1f2c4c
1 parent 12690d3
commit f1f2c4c
Show file tree

Hide file tree

Showing 7 changed files with 22 additions and 14 deletions.
diff --git a/README.md b/README.md
@@ -274,6 +274,7 @@ List of command-line flags
 |`--cfg-cache`                         | ExLlama_HF: Create an additional cache for CFG negative prompts. Necessary to use CFG with that loader, but not necessary for CFG with base ExLlama. |
 |`--no_flash_attn`                     | Force flash-attention to not be used. |
 |`--cache_8bit`                        | Use 8-bit cache to save VRAM. |
+|`--num_experts_per_token NUM_EXPERTS_PER_TOKEN` |  Number of experts to use for generation. Applies to MoE models like Mixtral. |
 
 #### AutoGPTQ
 

diff --git a/modules/exllamav2.py b/modules/exllamav2.py
@@ -48,6 +48,7 @@ def from_pretrained(self, path_to_model):
         config.scale_pos_emb = shared.args.compress_pos_emb
         config.scale_alpha_value = shared.args.alpha_value
         config.no_flash_attn = shared.args.no_flash_attn
+        config.num_experts_per_token = int(shared.args.num_experts_per_token)
 
         model = ExLlamaV2(config)
 

diff --git a/modules/exllamav2_hf.py b/modules/exllamav2_hf.py
@@ -165,5 +165,6 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
         config.scale_pos_emb = shared.args.compress_pos_emb
         config.scale_alpha_value = shared.args.alpha_value
         config.no_flash_attn = shared.args.no_flash_attn
+        config.num_experts_per_token = int(shared.args.num_experts_per_token)
 
         return Exllamav2HF(config)
diff --git a/modules/loaders.py b/modules/loaders.py
@@ -65,24 +65,25 @@
         'logits_all',
         'llamacpp_HF_info',
     ],
-    'ExLlama_HF': [
+    'ExLlamav2_HF': [
         'gpu_split',
         'max_seq_len',
+        'cfg_cache',
+        'no_flash_attn',
+        'num_experts_per_token',
+        'cache_8bit',
         'alpha_value',
-        'rope_freq_base',
         'compress_pos_emb',
-        'cfg_cache',
         'trust_remote_code',
         'no_use_fast',
     ],
-    'ExLlamav2_HF': [
+    'ExLlama_HF': [
         'gpu_split',
         'max_seq_len',
-        'cfg_cache',
-        'no_flash_attn',
-        'cache_8bit',
         'alpha_value',
+        'rope_freq_base',
         'compress_pos_emb',
+        'cfg_cache',
         'trust_remote_code',
         'no_use_fast',
     ],
@@ -123,22 +124,23 @@
         'no_use_fast',
         'gptq_for_llama_info',
     ],
-    'ExLlama': [
+    'ExLlamav2': [
         'gpu_split',
         'max_seq_len',
+        'no_flash_attn',
+        'num_experts_per_token',
+        'cache_8bit',
         'alpha_value',
-        'rope_freq_base',
         'compress_pos_emb',
-        'exllama_info',
+        'exllamav2_info',
     ],
-    'ExLlamav2': [
+    'ExLlama': [
         'gpu_split',
         'max_seq_len',
-        'no_flash_attn',
-        'cache_8bit',
         'alpha_value',
+        'rope_freq_base',
         'compress_pos_emb',
-        'exllamav2_info',
+        'exllama_info',
     ],
     'ctransformers': [
         'n_ctx',

diff --git a/modules/shared.py b/modules/shared.py
@@ -125,6 +125,7 @@
 parser.add_argument('--cfg-cache', action='store_true', help='ExLlama_HF: Create an additional cache for CFG negative prompts. Necessary to use CFG with that loader, but not necessary for CFG with base ExLlama.')
 parser.add_argument('--no_flash_attn', action='store_true', help='Force flash-attention to not be used.')
 parser.add_argument('--cache_8bit', action='store_true', help='Use 8-bit cache to save VRAM.')
+parser.add_argument('--num_experts_per_token', type=int, default=2, help='Number of experts to use for generation. Applies to MoE models like Mixtral.')
 
 # AutoGPTQ
 parser.add_argument('--triton', action='store_true', help='Use triton.')

diff --git a/modules/ui.py b/modules/ui.py
@@ -73,6 +73,7 @@ def list_model_elements():
         'disable_exllamav2',
         'cfg_cache',
         'no_flash_attn',
+        'num_experts_per_token',
         'cache_8bit',
         'threads',
         'threads_batch',

diff --git a/modules/ui_model_menu.py b/modules/ui_model_menu.py
@@ -129,6 +129,7 @@ def create_ui():
                             shared.gradio['no_flash_attn'] = gr.Checkbox(label="no_flash_attn", value=shared.args.no_flash_attn, info='Force flash-attention to not be used.')
                             shared.gradio['cache_8bit'] = gr.Checkbox(label="cache_8bit", value=shared.args.cache_8bit, info='Use 8-bit cache to save VRAM.')
                             shared.gradio['no_use_fast'] = gr.Checkbox(label="no_use_fast", value=shared.args.no_use_fast, info='Set use_fast=False while loading the tokenizer.')
+                            shared.gradio['num_experts_per_token'] = gr.Number(label="Number of experts per token", value=shared.args.num_experts_per_token, info='Only applies to MoE models like Mixtral.')
                             shared.gradio['gptq_for_llama_info'] = gr.Markdown('Legacy loader for compatibility with older GPUs. ExLlama_HF or AutoGPTQ are preferred for GPTQ models when supported.')
                             shared.gradio['exllama_info'] = gr.Markdown("ExLlama_HF is recommended over ExLlama for better integration with extensions and more consistent sampling behavior across loaders.")
                             shared.gradio['exllamav2_info'] = gr.Markdown("ExLlamav2_HF is recommended over ExLlamav2 for better integration with extensions and more consistent sampling behavior across loaders.")