Support studio (#300)

tastelikefeet · web-flow · commit 451194495637 · 2024-01-09T18:35:20.000+08:00
diff --git a/examples/pytorch/llm/scripts/llama2_13b_chat/longlora_ddp_ds/infer.sh b/examples/pytorch/llm/scripts/llama2_13b_chat/longlora_ddp_ds/infer.sh
@@ -0,0 +1,13 @@
+# Experimental environment: A100
+PYTHONPATH=../../.. \
+CUDA_VISIBLE_DEVICES=0 \
+python llm_infer.py \
+    --ckpt_dir "output/llama2-13b-chat/vx_xxx/checkpoint-xxx" \
+    --load_dataset_config true \
+    --max_length 4096 \
+    --max_new_tokens 2048 \
+    --temperature 0.1 \
+    --top_p 0.7 \
+    --repetition_penalty 1.05 \
+    --do_sample true \
+    --merge_lora_and_save false \
diff --git a/examples/pytorch/llm/scripts/llama2_13b_chat/longlora_ddp_ds/sft.sh b/examples/pytorch/llm/scripts/llama2_13b_chat/longlora_ddp_ds/sft.sh
@@ -0,0 +1,44 @@
+# Experimental environment: 2 * A100
+# 2 * 37GB GPU memory
+nproc_per_node=2
+
+PYTHONPATH=../../.. \
+CUDA_VISIBLE_DEVICES=0,1 \
+torchrun \
+    --nproc_per_node=$nproc_per_node \
+    --master_port 29500 \
+    llm_sft.py \
+    --model_id_or_path modelscope/Llama-2-13b-chat-ms \
+    --model_revision master \
+    --sft_type longlora \
+    --tuner_backend swift \
+    --template_type llama \
+    --dtype AUTO \
+    --output_dir output \
+    --ddp_backend nccl \
+    --dataset leetcode-python-en \
+    --train_dataset_sample -1 \
+    --num_train_epochs 1 \
+    --max_length 4096 \
+    --check_dataset_strategy warning \
+    --lora_rank 8 \
+    --lora_alpha 32 \
+    --lora_dropout_p 0.05 \
+    --lora_target_modules ALL \
+    --gradient_checkpointing true \
+    --batch_size 1 \
+    --weight_decay 0.01 \
+    --learning_rate 1e-4 \
+    --gradient_accumulation_steps $(expr 16 / $nproc_per_node) \
+    --max_grad_norm 0.5 \
+    --warmup_ratio 0.03 \
+    --eval_steps 100 \
+    --save_steps 100 \
+    --save_total_limit 2 \
+    --logging_steps 10 \
+    --push_to_hub false \
+    --hub_model_id llama2-13b-chat-longlora \
+    --hub_private_repo true \
+    --hub_token 'your-sdk-token' \
+    --deepspeed_config_path 'ds_config/zero2.json' \
+    --save_only_model true \
diff --git a/examples/pytorch/llm/scripts/qwen_7b_chat_int4/qalora/infer.sh b/examples/pytorch/llm/scripts/qwen_7b_chat_int4/qalora/infer.sh
@@ -0,0 +1,14 @@
+# Experimental environment: V100, A10, 3090
+PYTHONPATH=../../.. \
+CUDA_VISIBLE_DEVICES=0 \
+python llm_infer.py \
+    --ckpt_dir "output/qwen-7b-chat-int4/vx_xxx/checkpoint-xxx" \
+    --load_dataset_config true \
+    --max_length 4096 \
+    --use_flash_attn false \
+    --max_new_tokens 2048 \
+    --temperature 0.1 \
+    --top_p 0.7 \
+    --repetition_penalty 1.05 \
+    --do_sample true \
+    --merge_lora_and_save false \
diff --git a/examples/pytorch/llm/scripts/qwen_7b_chat_int4/qalora/sft.sh b/examples/pytorch/llm/scripts/qwen_7b_chat_int4/qalora/sft.sh
@@ -0,0 +1,37 @@
+# Experimental environment: V100, A10, 3090
+# 14GB GPU memory
+PYTHONPATH=../../.. \
+CUDA_VISIBLE_DEVICES=0 \
+python llm_sft.py \
+    --model_id_or_path qwen/Qwen-7B-Chat-Int4 \
+    --model_revision master \
+    --sft_type qalora \
+    --tuner_backend swift \
+    --template_type qwen \
+    --dtype fp16 \
+    --output_dir output \
+    --dataset leetcode-python-en \
+    --train_dataset_sample -1 \
+    --num_train_epochs 1 \
+    --max_length 4096 \
+    --check_dataset_strategy warning \
+    --lora_rank 8 \
+    --lora_alpha 32 \
+    --lora_dropout_p 0.05 \
+    --lora_target_modules ALL \
+    --gradient_checkpointing true \
+    --batch_size 1 \
+    --weight_decay 0.01 \
+    --learning_rate 1e-4 \
+    --gradient_accumulation_steps 16 \
+    --max_grad_norm 0.5 \
+    --warmup_ratio 0.03 \
+    --eval_steps 100 \
+    --save_steps 100 \
+    --save_total_limit 2 \
+    --logging_steps 10 \
+    --use_flash_attn false \
+    --push_to_hub false \
+    --hub_model_id qwen-7b-chat-int4-qalora \
+    --hub_private_repo true \
+    --hub_token 'your-sdk-token' \
diff --git a/swift/tuners/base.py b/swift/tuners/base.py
@@ -461,7 +461,7 @@ def set_active_adapters(self,
                             adapter_names: Union[List[str], str],
                             offload=None):
         if not adapter_names:
-            return
+            adapter_names = []
 
         if isinstance(adapter_names, str):
             adapter_names = [adapter_names]
diff --git a/swift/ui/app.py b/swift/ui/app.py
@@ -34,4 +34,5 @@ def run_ui():
             LLMTrain.build_ui(LLMTrain)
             LLMInfer.build_ui(LLMInfer)
 
-    app.queue().launch(height=800, share=False)
+    app.queue().launch(
+        height=800, share=bool(os.environ.get('WEBUI_SHARE', '0')))
diff --git a/swift/ui/base.py b/swift/ui/base.py
@@ -7,6 +7,8 @@
 from gradio import (Accordion, Button, Checkbox, Dropdown, Slider, Tab,
                     TabItem, Textbox)
 
+from swift.llm.utils.model import MODEL_MAPPING, ModelType
+
 all_langs = ['zh', 'en']
 builder: Type['BaseUI'] = None
 base_builder: Type['BaseUI'] = None
@@ -168,3 +170,8 @@ def get_default_value_from_dataclass(dataclass):
             else:
                 default_dict[f.name] = None
         return default_dict
+
+    @staticmethod
+    def get_custom_name_list():
+        return list(
+            set(MODEL_MAPPING.keys()) - set(ModelType.get_model_name_list()))
diff --git a/swift/ui/llm_infer/llm_infer.py b/swift/ui/llm_infer/llm_infer.py
@@ -1,6 +1,5 @@
 import os
 import re
-from dataclasses import fields
 from typing import Type
 
 import gradio as gr
@@ -138,7 +137,6 @@ def reset_memory(cls):
 
     @classmethod
     def prepare_checkpoint(cls, *args):
-        global model, tokenizer, template
         torch.cuda.empty_cache()
         infer_args = cls.get_default_value_from_dataclass(InferArguments)
         kwargs = {}
@@ -201,6 +199,8 @@ def generate_chat(cls, model_and_template, template_type, prompt: str,
             gr.Warning(cls.locale('generate_alert', cls.lang)['value'])
             return '', None
         model, template = model_and_template
+        if os.environ.get('MODELSCOPE_ENVIRONMENT') == 'studio':
+            model.cuda()
         if not template_type.endswith('generation'):
             old_history, history = limit_history_length(
                 template, prompt, history, int(max_new_tokens))
@@ -211,3 +211,5 @@ def generate_chat(cls, model_and_template, template_type, prompt: str,
         for _, history in gen:
             total_history = old_history + history
             yield '', total_history
+        if os.environ.get('MODELSCOPE_ENVIRONMENT') == 'studio':
+            model.cpu()
diff --git a/swift/ui/llm_infer/model.py b/swift/ui/llm_infer/model.py
@@ -86,7 +86,7 @@ def do_build_ui(cls, base_tab: Type['BaseUI']):
             model_type = gr.Dropdown(
                 elem_id='model_type',
                 choices=[base_tab.locale('checkpoint', cls.lang)['value']]
-                + ModelType.get_model_name_list(),
+                + ModelType.get_model_name_list() + cls.get_custom_name_list(),
                 value=base_tab.locale('checkpoint', cls.lang)['value'],
                 scale=20)
             model_id_or_path = gr.Textbox(
diff --git a/swift/ui/llm_train/llm_train.py b/swift/ui/llm_train/llm_train.py
@@ -1,6 +1,8 @@
+import collections
 import os
 import sys
 import time
+from subprocess import PIPE, STDOUT, Popen
 from typing import Dict, Type
 
 import gradio as gr
@@ -191,16 +193,31 @@ def do_build_ui(cls, base_tab: Type['BaseUI']):
                 Quantization.build_ui(base_tab)
                 SelfCog.build_ui(base_tab)
                 Advanced.build_ui(base_tab)
-                submit.click(
-                    cls.train, [
-                        value for value in cls.elements().values()
-                        if not isinstance(value, (Tab, Accordion))
-                    ], [
-                        cls.element('running_cmd'),
-                        cls.element('logging_dir'),
-                        cls.element('runtime_tab')
-                    ],
-                    show_progress=True)
+                if os.environ.get('MODELSCOPE_ENVIRONMENT') == 'studio':
+                    submit.click(
+                        cls.update_runtime, [],
+                        [cls.element('runtime_tab'),
+                         cls.element('log')]).then(
+                             cls.train_studio, [
+                                 value for value in cls.elements().values()
+                                 if not isinstance(value, (Tab, Accordion))
+                             ], [cls.element('log')],
+                             queue=True)
+                else:
+                    submit.click(
+                        cls.train_local, [
+                            value for value in cls.elements().values()
+                            if not isinstance(value, (Tab, Accordion))
+                        ], [
+                            cls.element('running_cmd'),
+                            cls.element('logging_dir'),
+                            cls.element('runtime_tab'),
+                        ],
+                        queue=True)
+
+    @classmethod
+    def update_runtime(cls):
+        return gr.update(visible=True), gr.update(visible=True)
 
     @classmethod
     def train(cls, *args):
@@ -239,7 +256,8 @@ def train(cls, *args):
                 params += f'--{e} {kwargs[e]} '
             else:
                 params += f'--{e} "{kwargs[e]}" '
-        params += '--add_output_dir_suffix False '
+        params += f'--add_output_dir_suffix False --output_dir {sft_args.output_dir} ' \
+                  f'--logging_dir {sft_args.logging_dir}'
         for key, param in more_params.items():
             params += f'--{key} "{param}" '
         ddp_param = ''
@@ -260,9 +278,30 @@ def train(cls, *args):
             if ddp_param:
                 ddp_param = f'set {ddp_param} && '
             run_command = f'{cuda_param}{ddp_param}start /b swift sft {params} > {log_file} 2>&1'
+        elif os.environ.get('MODELSCOPE_ENVIRONMENT') == 'studio':
+            run_command = f'{cuda_param} {ddp_param} swift sft {params}'
         else:
             run_command = f'{cuda_param} {ddp_param} nohup swift sft {params} > {log_file} 2>&1 &'
         logger.info(f'Run training: {run_command}')
+        return run_command, sft_args, other_kwargs
+
+    @classmethod
+    def train_studio(cls, *args):
+        run_command, sft_args, other_kwargs = cls.train(*args)
+        if os.environ.get('MODELSCOPE_ENVIRONMENT') == 'studio':
+            lines = collections.deque(
+                maxlen=int(os.environ.get('MAX_LOG_LINES', 50)))
+            process = Popen(
+                run_command, shell=True, stdout=PIPE, stderr=STDOUT)
+            with process.stdout:
+                for line in iter(process.stdout.readline, b''):
+                    line = line.decode('utf-8')
+                    lines.append(line)
+                    yield '\n'.join(lines)
+
+    @classmethod
+    def train_local(cls, *args):
+        run_command, sft_args, other_kwargs = cls.train(*args)
         if not other_kwargs['dry_run']:
             os.makedirs(sft_args.logging_dir, exist_ok=True)
             os.system(run_command)
diff --git a/swift/ui/llm_train/model.py b/swift/ui/llm_train/model.py
@@ -58,7 +58,8 @@ def do_build_ui(cls, base_tab: Type['BaseUI']):
         with gr.Row():
             model_type = gr.Dropdown(
                 elem_id='model_type',
-                choices=ModelType.get_model_name_list(),
+                choices=ModelType.get_model_name_list()
+                + cls.get_custom_name_list(),
                 scale=20)
             model_id_or_path = gr.Textbox(
                 elem_id='model_id_or_path',
diff --git a/swift/ui/llm_train/runtime.py b/swift/ui/llm_train/runtime.py
@@ -1,12 +1,18 @@
+import collections
 import os.path
+import time
 import webbrowser
 from typing import Dict, List, Tuple, Type
 
 import gradio as gr
+import psutil
 from transformers import is_tensorboard_available
 
 from swift.ui.base import BaseUI
 from swift.ui.llm_train.utils import close_loop, run_command_in_subprocess
+from swift.utils import get_logger
+
+logger = get_logger()
 
 
 class Runtime(BaseUI):
@@ -108,12 +114,15 @@ def do_build_ui(cls, base_tab: Type['BaseUI']):
                         max_lines=1)
                     gr.Button(elem_id='start_tb', scale=2, variant='primary')
                     gr.Button(elem_id='close_tb', scale=2)
+                with gr.Row():
+                    gr.Textbox(elem_id='log', lines=6, visible=False)
 
                 base_tab.element('show_log').click(
-                    Runtime.show_log,
-                    [base_tab.element('logging_dir')],
-                    [],
-                )
+                    Runtime.update_log, [], [cls.element('log')]).then(
+                        Runtime.wait, [base_tab.element('logging_dir')],
+                        [cls.element('log')],
+                        show_progress=True,
+                        queue=True)
 
                 base_tab.element('start_tb').click(
                     Runtime.start_tb,
@@ -127,6 +136,52 @@ def do_build_ui(cls, base_tab: Type['BaseUI']):
                     [],
                 )
 
+    @classmethod
+    def update_log(cls):
+        return gr.update(visible=True)
+
+    @classmethod
+    def wait(cls, logging_dir):
+        log_file = os.path.join(logging_dir, 'run.log')
+        offset = 0
+        latest_data = ''
+        lines = collections.deque(
+            maxlen=int(os.environ.get('MAX_LOG_LINES', 50)))
+        while True:
+            try:
+                with open(log_file) as input:
+                    input.seek(offset)
+                    fail_cnt = 0
+                    while True:
+                        latest_data += input.read()
+                        offset = input.tell()
+                        if not latest_data:
+                            time.sleep(0.5)
+                            fail_cnt += 1
+                            if fail_cnt > 5:
+                                break
+
+                        if '\n' not in latest_data:
+                            continue
+                        latest_lines = latest_data.split('\n')
+                        if latest_data[-1] != '\n':
+                            latest_data = latest_lines[-1]
+                            latest_lines = latest_lines[:-1]
+                        else:
+                            latest_data = ''
+                        lines.extend(latest_lines)
+                        yield '\n'.join(lines)
+            except IOError:
+                pass
+
+            process_name = 'swift'
+            process_find = False
+            for proc in psutil.process_iter():
+                if proc.name() == process_name:
+                    process_find = proc.pid
+            if not process_find:
+                break
+
     @classmethod
     def show_log(cls, logging_dir):
         webbrowser.open(
@@ -152,6 +207,8 @@ def start_tb(cls, logging_dir):
                 line = line[line.index('http://localhost:'):]
                 localhost_addr = line[:line.index(' ')]
         cls.handlers[logging_dir] = (handler, localhost_addr)
+        logger.info('===========Tensorboard Log============')
+        logger.info('\n'.join(lines))
         webbrowser.open(localhost_addr, new=2)
         return localhost_addr
 
