Fix FSDP; Add training percentage to jsonl logging; Add a web-ui component (#1381)

Author: tastelikefeet

swift/llm/sft.py

@@ -18,9 +18,8 @@
                          plot_images, preprocess_logits_for_metrics, seed_everything, show_layers, use_torchacc)
 from .accelerator import ta_accelerate
 from .tuner import prepare_model
-from .utils import (TEMPLATE_MAPPING, LazyLLMDataset, SftArguments, Template, dataset_map, get_dataset,
-                    get_model_tokenizer, get_template, get_time_info, print_example, set_generation_config,
-                    sort_by_max_length, stat_dataset)
+from .utils import (LazyLLMDataset, SftArguments, Template, dataset_map, get_dataset, get_model_tokenizer, get_template,
+                    get_time_info, print_example, set_generation_config, sort_by_max_length, stat_dataset)
 
 logger = get_logger()
 
@@ -42,7 +41,7 @@ def llm_sft(args: SftArguments) -> Dict[str, Union[str, Any]]:
             torch.cuda.set_per_process_memory_fraction(max(min(args.gpu_memory_fraction, 1.0), 0.01), device=device_id)
 
     # Loading Model and Tokenizer
-    if is_deepspeed_zero3_enabled():
+    if is_deepspeed_zero3_enabled() or os.environ.get('ACCELERATE_USE_FSDP', 'False') == 'true':
         model_kwargs = {'device_map': None}
     elif is_torch_npu_available():
         model_kwargs = {'device_map': local_rank if local_rank >= 0 else 0}

swift/trainers/callback.py

@@ -8,6 +8,7 @@
 from transformers.trainer_utils import IntervalStrategy, has_length, speed_metrics
 
 from swift.utils import append_to_jsonl, is_pai_training_job, use_torchacc
+from ..utils.utils import format_time
 from .arguments import TrainingArguments
 
 
@@ -17,6 +18,7 @@ def on_train_begin(self, args, state, control, **kwargs):
         if state.is_local_process_zero:
             self.training_bar = tqdm(desc='Train', total=state.max_steps, dynamic_ncols=True)
         self.current_step = 0
+        self.start_time = time.time()
         if use_torchacc():
             self.warmup_start_time = 0
             self.warmup_metric = None
@@ -33,7 +35,14 @@ def on_prediction_step(self, args, state: TrainerState, control, eval_dataloader
             self.prediction_bar.update()
 
     def on_log(self, args: TrainingArguments, state: TrainerState, control, logs=None, **kwargs):
-        logs['global_step'] = state.global_step
+        logs['steps[global_step/max_steps]'] = f'{state.global_step}/{state.max_steps}'
+        train_percentage = state.global_step / state.max_steps if state.max_steps else 0.
+        logs['percentage'] = f'{train_percentage * 100:.2f}%'
+        elapsed = time.time() - self.start_time
+        elapsed = max(0., elapsed)
+        logs['elapsed_time'] = format_time(elapsed)
+        logs['remaining_time'] = format_time(elapsed / train_percentage - elapsed)
+
         if use_torchacc():
             if state.global_step >= self.metric_warmup_step and self.warmup_start_time == 0:
                 self.warmup_start_time = time.time()

swift/ui/base.py

@@ -1,10 +1,10 @@
 import os
 import typing
 from dataclasses import fields
-from functools import partial, wraps
+from functools import wraps
 from typing import Any, Dict, List, OrderedDict, Type
 
-from gradio import Accordion, Button, Checkbox, Dropdown, Slider, Tab, TabItem, Textbox
+from gradio import Accordion, Audio, Button, Checkbox, Dropdown, File, Image, Slider, Tab, TabItem, Textbox, Video
 
 from swift.llm.utils.model import MODEL_MAPPING, ModelType
 
@@ -69,6 +69,10 @@ def wrapper(*args, **kwargs):
 TabItem.__init__ = update_data(TabItem.__init__)
 Accordion.__init__ = update_data(Accordion.__init__)
 Button.__init__ = update_data(Button.__init__)
+File.__init__ = update_data(File.__init__)
+Image.__init__ = update_data(Image.__init__)
+Video.__init__ = update_data(Video.__init__)
+Audio.__init__ = update_data(Audio.__init__)
 
 
 class BaseUI:

swift/ui/llm_eval/llm_eval.py

@@ -126,7 +126,7 @@ def eval(cls, *args):
                 elif isinstance(value, str) and re.fullmatch(cls.bool_regex, value):
                     value = True if value.lower() == 'true' else False
                 kwargs[key] = value if not isinstance(value, list) else ' '.join(value)
-                kwargs_is_list[key] = isinstance(value, list)
+                kwargs_is_list[key] = isinstance(value, list) or getattr(cls.element(key), 'is_list', False)
             else:
                 other_kwargs[key] = value
             if key == 'more_params' and value:

swift/ui/llm_export/llm_export.py

@@ -124,7 +124,7 @@ def export(cls, *args):
                 elif isinstance(value, str) and re.fullmatch(cls.bool_regex, value):
                     value = True if value.lower() == 'true' else False
                 kwargs[key] = value if not isinstance(value, list) else ' '.join(value)
-                kwargs_is_list[key] = isinstance(value, list)
+                kwargs_is_list[key] = isinstance(value, list) or getattr(cls.element(key), 'is_list', False)
             else:
                 other_kwargs[key] = value
             if key == 'more_params' and value:

swift/ui/llm_infer/llm_infer.py

@@ -2,7 +2,6 @@
 import re
 import sys
 import time
-from copy import copy
 from datetime import datetime
 from functools import partial
 from typing import Type
@@ -14,7 +13,7 @@
 from modelscope import GenerationConfig, snapshot_download
 
 from swift.llm import (TEMPLATE_MAPPING, DeployArguments, InferArguments, XRequestConfig, inference_client,
-                       inference_stream, limit_history_length, prepare_model_template)
+                       inference_stream, prepare_model_template)
 from swift.ui.base import BaseUI
 from swift.ui.llm_infer.model import Model
 from swift.ui.llm_infer.runtime import Runtime
@@ -69,6 +68,16 @@ class LLMInfer(BaseUI):
                 'en': 'Chat bot'
             },
         },
+        'infer_model_type': {
+            'label': {
+                'zh': 'Lora模块',
+                'en': 'Lora module'
+            },
+            'info': {
+                'zh': '发送给server端哪个LoRA，默认为`default-lora`',
+                'en': 'Which LoRA to use on server, default value is `default-lora`'
+            }
+        },
         'prompt': {
             'label': {
                 'zh': '请输入：',
@@ -116,12 +125,14 @@ def do_build_ui(cls, base_tab: Type['BaseUI']):
                 history = gr.State([])
                 Model.build_ui(base_tab)
                 Runtime.build_ui(base_tab)
-                gr.Dropdown(
-                    elem_id='gpu_id',
-                    multiselect=True,
-                    choices=[str(i) for i in range(gpu_count)] + ['cpu'],
-                    value=default_device,
-                    scale=8)
+                with gr.Row():
+                    gr.Dropdown(
+                        elem_id='gpu_id',
+                        multiselect=True,
+                        choices=[str(i) for i in range(gpu_count)] + ['cpu'],
+                        value=default_device,
+                        scale=8)
+                    infer_model_type = gr.Textbox(elem_id='infer_model_type', scale=4)
                 chatbot = gr.Chatbot(elem_id='chatbot', elem_classes='control-height')
                 with gr.Row():
                     prompt = gr.Textbox(elem_id='prompt', lines=1, interactive=True)
@@ -172,7 +183,7 @@ def do_build_ui(cls, base_tab: Type['BaseUI']):
                         cls.send_message,
                         inputs=[
                             cls.element('running_tasks'), model_and_template,
-                            cls.element('template_type'), prompt, image, history,
+                            cls.element('template_type'), prompt, image, history, infer_model_type,
                             cls.element('system'),
                             cls.element('max_new_tokens'),
                             cls.element('temperature'),
@@ -217,7 +228,7 @@ def deploy(cls, *args):
                 elif isinstance(value, str) and re.fullmatch(cls.bool_regex, value):
                     value = True if value.lower() == 'true' else False
                 kwargs[key] = value if not isinstance(value, list) else ' '.join(value)
-                kwargs_is_list[key] = isinstance(value, list)
+                kwargs_is_list[key] = isinstance(value, list) or getattr(cls.element(key), 'is_list', False)
             else:
                 other_kwargs[key] = value
             if key == 'more_params' and value:
@@ -374,8 +385,8 @@ def agent_type(cls, response):
         return None
 
     @classmethod
-    def send_message(cls, running_task, model_and_template, template_type, prompt: str, image, history, system,
-                     max_new_tokens, temperature, top_k, top_p, repetition_penalty):
+    def send_message(cls, running_task, model_and_template, template_type, prompt: str, image, history,
+                     infer_model_type, system, max_new_tokens, temperature, top_k, top_p, repetition_penalty):
         if not model_and_template:
             gr.Warning(cls.locale('generate_alert', cls.lang)['value'])
             return '', None, None, []
@@ -393,7 +404,7 @@ def send_message(cls, running_task, model_and_template, template_type, prompt: s
         _, args = Runtime.parse_info_from_cmdline(running_task)
         model_type, template, sft_type = model_and_template
         if sft_type in ('lora', 'longlora') and not args.get('merge_lora'):
-            model_type = 'default-lora'
+            model_type = infer_model_type or 'default-lora'
         old_history, history = history or [], []
         request_config = XRequestConfig(
             temperature=temperature, top_k=top_k, top_p=top_p, repetition_penalty=repetition_penalty)

swift/ui/llm_infer/model.py

@@ -80,6 +80,16 @@ class Model(BaseUI):
                 'en': 'Only available when sft_type=lora'
             }
         },
+        'lora_modules': {
+            'label': {
+                'zh': '外部lora模块',
+                'en': 'More lora modules'
+            },
+            'info': {
+                'zh': '空格分割的name=/path1/path2键值对',
+                'en': 'name=/path1/path2 split by blanks'
+            }
+        },
         'more_params': {
             'label': {
                 'zh': '更多参数',
@@ -117,6 +127,7 @@ def do_build_ui(cls, base_tab: Type['BaseUI']):
             system = gr.Textbox(elem_id='system', lines=4, scale=20)
         Generate.build_ui(base_tab)
         with gr.Row():
+            gr.Textbox(elem_id='lora_modules', lines=1, is_list=True, scale=40)
             gr.Textbox(elem_id='more_params', lines=1, scale=20)
             gr.Button(elem_id='load_checkpoint', scale=2, variant='primary')
 

swift/ui/llm_infer/runtime.py

@@ -13,6 +13,7 @@
 
 from swift.ui.base import BaseUI
 from swift.utils import get_logger
+from swift.utils.utils import format_time
 
 logger = get_logger()
 
@@ -211,23 +212,6 @@ def construct_running_task(proc):
         create_time = proc.create_time()
         create_time_formatted = datetime.fromtimestamp(create_time).strftime('%Y-%m-%d, %H:%M')
 
-        def format_time(seconds):
-            days = int(seconds // (24 * 3600))
-            hours = int((seconds % (24 * 3600)) // 3600)
-            minutes = int((seconds % 3600) // 60)
-            seconds = int(seconds % 60)
-
-            if days > 0:
-                time_str = f'{days}d {hours}h {minutes}m {seconds}s'
-            elif hours > 0:
-                time_str = f'{hours}h {minutes}m {seconds}s'
-            elif minutes > 0:
-                time_str = f'{minutes}m {seconds}s'
-            else:
-                time_str = f'{seconds}s'
-
-            return time_str
-
         return f'pid:{pid}/create:{create_time_formatted}' \
                f'/running:{format_time(ts - create_time)}/cmd:{" ".join(proc.cmdline())}'
 

swift/ui/llm_train/runtime.py

@@ -16,6 +16,7 @@
 from swift.ui.base import BaseUI
 from swift.ui.llm_train.utils import close_loop, run_command_in_subprocess
 from swift.utils import TB_COLOR, TB_COLOR_SMOOTH, get_logger, read_tensorboard_file, tensorboard_smoothing
+from swift.utils.utils import format_time
 
 logger = get_logger()
 
@@ -423,23 +424,6 @@ def construct_running_task(proc):
         create_time = proc.create_time()
         create_time_formatted = datetime.fromtimestamp(create_time).strftime('%Y-%m-%d, %H:%M')
 
-        def format_time(seconds):
-            days = int(seconds // (24 * 3600))
-            hours = int((seconds % (24 * 3600)) // 3600)
-            minutes = int((seconds % 3600) // 60)
-            seconds = int(seconds % 60)
-
-            if days > 0:
-                time_str = f'{days}d {hours}h {minutes}m {seconds}s'
-            elif hours > 0:
-                time_str = f'{hours}h {minutes}m {seconds}s'
-            elif minutes > 0:
-                time_str = f'{minutes}m {seconds}s'
-            else:
-                time_str = f'{seconds}s'
-
-            return time_str
-
         return f'pid:{pid}/create:{create_time_formatted}' \
                f'/running:{format_time(ts-create_time)}/cmd:{" ".join(proc.cmdline())}'
 

swift/utils/utils.py

@@ -63,6 +63,24 @@ def _get_version(work_dir: str) -> int:
     return max(v_list) + 1
 
 
+def format_time(seconds):
+    days = int(seconds // (24 * 3600))
+    hours = int((seconds % (24 * 3600)) // 3600)
+    minutes = int((seconds % 3600) // 60)
+    seconds = int(seconds % 60)
+
+    if days > 0:
+        time_str = f'{days}d {hours}h {minutes}m {seconds}s'
+    elif hours > 0:
+        time_str = f'{hours}h {minutes}m {seconds}s'
+    elif minutes > 0:
+        time_str = f'{minutes}m {seconds}s'
+    else:
+        time_str = f'{seconds}s'
+
+    return time_str
+
+
 def seed_everything(seed: Optional[int] = None, full_determinism: bool = False, *, verbose: bool = True) -> int:
 
     if seed is None: