From 93c250b9b6b7d625c9cd698f99400d4fe650779d Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Tue, 1 Oct 2024 11:16:15 -0700
Subject: [PATCH] Add a UI element for enable_tp

---
 modules/loaders.py       | 2 ++
 modules/ui.py            | 1 +
 modules/ui_model_menu.py | 1 +
 3 files changed, 4 insertions(+)

diff --git a/modules/loaders.py b/modules/loaders.py
index 16a3106e..deee00a7 100644
--- a/modules/loaders.py
+++ b/modules/loaders.py
@@ -90,6 +90,7 @@ loaders_and_params = OrderedDict({
         'cache_8bit',
         'cache_4bit',
         'autosplit',
+        'enable_tp',
         'alpha_value',
         'compress_pos_emb',
         'trust_remote_code',
@@ -105,6 +106,7 @@ loaders_and_params = OrderedDict({
         'cache_8bit',
         'cache_4bit',
         'autosplit',
+        'enable_tp',
         'alpha_value',
         'compress_pos_emb',
         'exllamav2_info',
diff --git a/modules/ui.py b/modules/ui.py
index 76b1c009..c07beeb4 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -90,6 +90,7 @@ def list_model_elements():
         'cache_8bit',
         'cache_4bit',
         'autosplit',
+        'enable_tp',
         'threads',
         'threads_batch',
         'n_batch',
diff --git a/modules/ui_model_menu.py b/modules/ui_model_menu.py
index 1883fdca..f87b680a 100644
--- a/modules/ui_model_menu.py
+++ b/modules/ui_model_menu.py
@@ -136,6 +136,7 @@ def create_ui():
                             shared.gradio['disk'] = gr.Checkbox(label="disk", value=shared.args.disk)
                             shared.gradio['bf16'] = gr.Checkbox(label="bf16", value=shared.args.bf16)
                             shared.gradio['autosplit'] = gr.Checkbox(label="autosplit", value=shared.args.autosplit, info='Automatically split the model tensors across the available GPUs.')
+                            shared.gradio['enable_tp'] = gr.Checkbox(label="enable_tp", value=shared.args.enable_tp, info='Enable Tensor Parallelism (TP).')
                             shared.gradio['no_flash_attn'] = gr.Checkbox(label="no_flash_attn", value=shared.args.no_flash_attn)
                             shared.gradio['no_xformers'] = gr.Checkbox(label="no_xformers", value=shared.args.no_xformers)
                             shared.gradio['no_sdpa'] = gr.Checkbox(label="no_sdpa", value=shared.args.no_sdpa)