add export_device in webui #3333

Former-commit-id: 30ebd3652809d73941e0a5e4a8be11d989faf98d

src/llmtuner/webui/components/export.py

@@ -1,5 +1,6 @@
 from typing import TYPE_CHECKING, Dict, Generator, List
 
+from ...extras.misc import torch_gc
 from ...extras.packages import is_gradio_available
 from ...train import export_model
 from ..common import get_save_dir
@@ -26,9 +27,10 @@ def save_model(
     adapter_path: List[str],
     finetuning_type: str,
     template: str,
-    max_shard_size: int,
+    export_size: int,
     export_quantization_bit: int,
     export_quantization_dataset: str,
+    export_device: str,
     export_legacy_format: bool,
     export_dir: str,
     export_hub_model_id: str,
@@ -44,6 +46,8 @@ def save_model(
         error = ALERTS["err_no_dataset"][lang]
     elif export_quantization_bit not in GPTQ_BITS and not adapter_path:
         error = ALERTS["err_no_adapter"][lang]
+    elif export_quantization_bit in GPTQ_BITS and adapter_path:
+        error = ALERTS["err_gptq_lora"][lang]
 
     if error:
         gr.Warning(error)
@@ -64,22 +68,25 @@ def save_model(
         template=template,
         export_dir=export_dir,
         export_hub_model_id=export_hub_model_id or None,
-        export_size=max_shard_size,
+        export_size=export_size,
         export_quantization_bit=int(export_quantization_bit) if export_quantization_bit in GPTQ_BITS else None,
         export_quantization_dataset=export_quantization_dataset,
+        export_device=export_device,
         export_legacy_format=export_legacy_format,
     )
 
     yield ALERTS["info_exporting"][lang]
     export_model(args)
+    torch_gc()
     yield ALERTS["info_exported"][lang]
 
 
 def create_export_tab(engine: "Engine") -> Dict[str, "Component"]:
     with gr.Row():
-        max_shard_size = gr.Slider(value=1, minimum=1, maximum=100, step=1)
+        export_size = gr.Slider(value=1, minimum=1, maximum=100, step=1)
         export_quantization_bit = gr.Dropdown(choices=["none", "8", "4", "3", "2"], value="none")
         export_quantization_dataset = gr.Textbox(value="data/c4_demo.json")
+        export_device = gr.Radio(choices=["cpu", "cuda"], value="cpu")
         export_legacy_format = gr.Checkbox()
 
     with gr.Row():
@@ -98,9 +105,10 @@ def create_export_tab(engine: "Engine") -> Dict[str, "Component"]:
             engine.manager.get_elem_by_id("top.adapter_path"),
             engine.manager.get_elem_by_id("top.finetuning_type"),
             engine.manager.get_elem_by_id("top.template"),
-            max_shard_size,
+            export_size,
             export_quantization_bit,
             export_quantization_dataset,
+            export_device,
             export_legacy_format,
             export_dir,
             export_hub_model_id,
@@ -109,9 +117,10 @@ def create_export_tab(engine: "Engine") -> Dict[str, "Component"]:
     )
 
     return dict(
-        max_shard_size=max_shard_size,
+        export_size=export_size,
         export_quantization_bit=export_quantization_bit,
         export_quantization_dataset=export_quantization_dataset,
+        export_device=export_device,
         export_legacy_format=export_legacy_format,
         export_dir=export_dir,
         export_hub_model_id=export_hub_model_id,