RLHF UI (modelscope#1182)

tastelikefeet · web-flow · commit a5f5000c52a3 · 2024-06-20T14:58:13.000+08:00
diff --git a/swift/llm/utils/argument.py b/swift/llm/utils/argument.py
@@ -1407,8 +1407,8 @@ class RLHFArguments(SftArguments):
     max_prompt_length: int = 1024
     beta: Optional[float] = None
     label_smoothing: float = 0.0
-    loss_type: Optional[Literal['sigmoid', 'hinge', 'ipo', 'kto_pair', 'robust', 'bco_pair', 'sppo_hard', 'nca_pair',
-                                'simpo', 'kto', 'bco']] = None
+    loss_type: Literal['sigmoid', 'hinge', 'ipo', 'kto_pair', 'robust', 'bco_pair', 'sppo_hard', 'nca_pair', 'simpo',
+                       'kto', 'bco'] = None
     sft_beta: float = 0.1
     simpo_gamma: float = 1.0  # reward margin hyperparameter in SimPO
     # KTO
diff --git a/swift/ui/base.py b/swift/ui/base.py
@@ -51,6 +51,7 @@ def wrapper(*args, **kwargs):
                 if argument and 'label' in kwargs:
                     kwargs['label'] = kwargs['label'] + f'({argument})'
 
+        kwargs['elem_classes'] = 'align'
         ret = fn(self, **kwargs)
         self.constructor_args.update(kwargs)
 
diff --git a/swift/ui/llm_train/dataset.py b/swift/ui/llm_train/dataset.py
@@ -113,7 +113,7 @@ def do_build_ui(cls, base_tab: Type['BaseUI']):
             gr.Textbox(elem_id='custom_val_dataset_path', is_list=True, scale=20)
         with gr.Row():
             gr.Slider(elem_id='dataset_test_ratio', minimum=0.0, maximum=1.0, step=0.05, scale=20)
-            gr.Slider(elem_id='max_length', minimum=32, maximum=8192, step=32, scale=20)
+            gr.Slider(elem_id='max_length', minimum=32, maximum=32768, step=32, scale=20)
             gr.Textbox(elem_id='train_dataset_sample', scale=20)
             gr.Textbox(elem_id='val_dataset_sample', scale=20)
             gr.Dropdown(elem_id='truncation_strategy', scale=20)
diff --git a/swift/ui/llm_train/llm_train.py b/swift/ui/llm_train/llm_train.py
@@ -12,7 +12,7 @@
 import torch
 from gradio import Accordion, Tab
 
-from swift.llm import SftArguments
+from swift.llm import RLHFArguments, SftArguments
 from swift.ui.base import BaseUI
 from swift.ui.llm_train.advanced import Advanced
 from swift.ui.llm_train.dataset import Dataset
@@ -23,6 +23,7 @@
 from swift.ui.llm_train.lora import LoRA
 from swift.ui.llm_train.model import Model
 from swift.ui.llm_train.quantization import Quantization
+from swift.ui.llm_train.rlhf import RLHF
 from swift.ui.llm_train.runtime import Runtime
 from swift.ui.llm_train.save import Save
 from swift.ui.llm_train.self_cog import SelfCog
@@ -53,6 +54,7 @@ class LLMTrain(BaseUI):
         Quantization,
         SelfCog,
         Advanced,
+        RLHF,
     ]
 
     locale_dict: Dict[str, Dict] = {
@@ -62,6 +64,16 @@ class LLMTrain(BaseUI):
                 'en': 'LLM Training',
             }
         },
+        'train_type': {
+            'label': {
+                'zh': '训练Stage',
+                'en': 'Train Stage'
+            },
+            'info': {
+                'zh': '请注意选择于此匹配的数据集，人类对齐配置在页面下方',
+                'en': 'Please choose matched dataset, RLHF settings is at the bottom of the page'
+            }
+        },
         'submit_alert': {
             'value': {
                 'zh':
@@ -185,9 +197,9 @@ class LLMTrain(BaseUI):
         },
     }
 
-    choice_dict = BaseUI.get_choices_from_dataclass(SftArguments)
-    default_dict = BaseUI.get_default_value_from_dataclass(SftArguments)
-    arguments = BaseUI.get_argument_names(SftArguments)
+    choice_dict = BaseUI.get_choices_from_dataclass(RLHFArguments)
+    default_dict = BaseUI.get_default_value_from_dataclass(RLHFArguments)
+    arguments = BaseUI.get_argument_names(RLHFArguments)
 
     @classmethod
     def do_build_ui(cls, base_tab: Type['BaseUI']):
@@ -200,17 +212,19 @@ def do_build_ui(cls, base_tab: Type['BaseUI']):
             with gr.Blocks():
                 Model.build_ui(base_tab)
                 Dataset.build_ui(base_tab)
-                Hyper.build_ui(base_tab)
-                Save.build_ui(base_tab)
-                Runtime.build_ui(base_tab)
                 with gr.Row():
-                    gr.Dropdown(elem_id='sft_type', scale=4)
-                    gr.Dropdown(elem_id='tuner_backend', scale=4)
-                    gr.Textbox(elem_id='sequence_parallel_size', scale=4)
+                    gr.Dropdown(elem_id='train_type', choices=['pretrain/sft', 'rlhf'], value='pretrain/sft', scale=3)
+                    gr.Dropdown(elem_id='sft_type', scale=2)
+                    gr.Dropdown(elem_id='tuner_backend', scale=2)
+                    gr.Textbox(elem_id='sequence_parallel_size', scale=3)
+                with gr.Row():
                     gr.Textbox(elem_id='seed', scale=4)
                     gr.Dropdown(elem_id='dtype', scale=4)
                     gr.Checkbox(elem_id='use_ddp', value=False, scale=4)
                     gr.Textbox(elem_id='ddp_num', value='2', scale=4)
+                Hyper.build_ui(base_tab)
+                Save.build_ui(base_tab)
+                Runtime.build_ui(base_tab)
                 with gr.Row():
                     gr.Dropdown(
                         elem_id='gpu_id',
@@ -230,6 +244,7 @@ def do_build_ui(cls, base_tab: Type['BaseUI']):
                 Lisa.build_ui(base_tab)
                 LlamaPro.build_ui(base_tab)
                 Quantization.build_ui(base_tab)
+                RLHF.build_ui(base_tab)
                 SelfCog.build_ui(base_tab)
                 Advanced.build_ui(base_tab)
 
@@ -273,26 +288,24 @@ def update_runtime(cls):
 
     @classmethod
     def train(cls, *args):
-        ignore_elements = ('model_type', 'logging_dir', 'more_params')
-        sft_args = cls.get_default_value_from_dataclass(SftArguments)
+        ignore_elements = ('model_type', 'logging_dir', 'more_params', 'train_type')
+        sft_args = cls.get_default_value_from_dataclass(RLHFArguments)
         kwargs = {}
         kwargs_is_list = {}
         other_kwargs = {}
         more_params = {}
         keys = [key for key, value in cls.elements().items() if not isinstance(value, (Tab, Accordion))]
         model_type = None
+        do_rlhf = False
         for key, value in zip(keys, args):
             compare_value = sft_args.get(key)
-            compare_value_arg = str(compare_value) if not isinstance(compare_value, (list, dict)) else compare_value
-            compare_value_ui = str(value) if not isinstance(value, (list, dict)) else value
-
             if isinstance(value, str) and re.fullmatch(cls.int_regex, value):
                 value = int(value)
             elif isinstance(value, str) and re.fullmatch(cls.float_regex, value):
                 value = float(value)
             elif isinstance(value, str) and re.fullmatch(cls.bool_regex, value):
                 value = True if value.lower() == 'true' else False
-            if key not in ignore_elements and key in sft_args and compare_value_ui != compare_value_arg and value:
+            if key not in ignore_elements and key in sft_args and compare_value != value and value:
                 kwargs[key] = value if not isinstance(value, list) else ' '.join(value)
                 kwargs_is_list[key] = isinstance(value, list) or getattr(cls.element(key), 'is_list', False)
             else:
@@ -303,14 +316,18 @@ def train(cls, *args):
             if key == 'model_type':
                 model_type = value
 
+            if key == 'train_type':
+                do_rlhf = value == 'rlhf'
+
         if os.path.exists(kwargs['model_id_or_path']):
             kwargs['model_type'] = model_type
 
         kwargs.update(more_params)
         if 'dataset' not in kwargs and 'custom_train_dataset_path' not in kwargs:
             raise gr.Error(cls.locale('dataset_alert', cls.lang)['value'])
 
-        sft_args = SftArguments(
+        cmd = 'rlhf' if do_rlhf else 'sft'
+        sft_args = RLHFArguments(
             **{
                 key: value.split(' ') if kwargs_is_list.get(key, False) and isinstance(value, str) else value
                 for key, value in kwargs.items()
@@ -323,7 +340,7 @@ def train(cls, *args):
             else:
                 params += f'--{e} "{kwargs[e]}" '
         params += f'--add_output_dir_suffix False --output_dir {sft_args.output_dir} ' \
-                  f'--logging_dir {sft_args.logging_dir} '
+                  f'--logging_dir {sft_args.logging_dir} --ignore_args_error True'
         ddp_param = ''
         devices = other_kwargs['gpu_id']
         devices = [d for d in devices if d]
@@ -344,9 +361,9 @@ def train(cls, *args):
                 ddp_param = f'set {ddp_param} && '
             run_command = f'{cuda_param}{ddp_param}start /b swift sft {params} > {log_file} 2>&1'
         elif cls.is_studio:
-            run_command = f'{cuda_param} {ddp_param} swift sft {params}'
+            run_command = f'{cuda_param} {ddp_param} swift {cmd} {params}'
         else:
-            run_command = f'{cuda_param} {ddp_param} nohup swift sft {params} > {log_file} 2>&1 &'
+            run_command = f'{cuda_param} {ddp_param} nohup swift {cmd} {params} > {log_file} 2>&1 &'
         logger.info(f'Run training: {run_command}')
         return run_command, sft_args, other_kwargs
 
diff --git a/swift/ui/llm_train/lora.py b/swift/ui/llm_train/lora.py
@@ -80,10 +80,10 @@ def do_build_ui(cls, base_tab: Type['BaseUI']):
         with gr.Accordion(elem_id='lora_tab', open=True):
             with gr.Blocks():
                 with gr.Row():
-                    lora_target_modules = gr.Textbox(elem_id='lora_target_modules', lines=1, scale=20, is_list=True)
+                    lora_target_modules = gr.Textbox(elem_id='lora_target_modules', lines=1, scale=5, is_list=True)
+                    gr.Slider(elem_id='lora_rank', value=32, minimum=1, maximum=512, step=8, scale=2)
+                    gr.Slider(elem_id='lora_alpha', value=8, minimum=1, maximum=512, step=8, scale=2)
                 with gr.Row():
-                    gr.Slider(elem_id='lora_rank', value=32, minimum=1, maximum=512, step=8)
-                    gr.Slider(elem_id='lora_alpha', value=8, minimum=1, maximum=512, step=8)
                     gr.Dropdown(elem_id='lora_dtype')
                     gr.Textbox(elem_id='lora_lr_ratio')
                     gr.Checkbox(elem_id='use_rslora')
diff --git a/swift/ui/llm_train/rlhf.py b/swift/ui/llm_train/rlhf.py
@@ -0,0 +1,149 @@
+from typing import Type
+
+import gradio as gr
+
+from swift.llm import MODEL_MAPPING, ModelType
+from swift.ui.base import BaseUI
+
+
+class RLHF(BaseUI):
+
+    group = 'llm_train'
+
+    locale_dict = {
+        'rlhf_tab': {
+            'label': {
+                'zh': '人类对齐参数设置',
+                'en': 'RLHF settings'
+            },
+        },
+        'rlhf_type': {
+            'label': {
+                'zh': '人类对齐算法类型',
+                'en': 'RLHF type'
+            },
+        },
+        'ref_model_type': {
+            'label': {
+                'zh': '选择ref模型',
+                'en': 'Select ref model'
+            },
+            'info': {
+                'zh': 'SWIFT已支持的模型名称',
+                'en': 'Base model supported by SWIFT'
+            }
+        },
+        'ref_model_id_or_path': {
+            'label': {
+                'zh': 'ref模型id或路径',
+                'en': 'Ref model id or path'
+            },
+            'info': {
+                'zh': '实际的模型id或路径',
+                'en': 'The actual model id or path'
+            }
+        },
+        'max_prompt_length': {
+            'label': {
+                'zh': 'prompt最大token长度',
+                'en': 'Max prompt length'
+            },
+        },
+        'beta': {
+            'label': {
+                'zh': 'KL正则项系数',
+                'en': 'KL regression ratio'
+            },
+        },
+        'loss_type': {
+            'label': {
+                'zh': 'Loss类型',
+                'en': 'Loss type'
+            },
+        },
+        'sft_beta': {
+            'label': {
+                'zh': 'DPO中混合sft交叉熵的系数',
+                'en': 'DPO Cross Entropy ratio'
+            },
+        },
+        'simpo_gamma': {
+            'label': {
+                'zh': 'SimPO reward margin',
+                'en': 'SimPO reward margin'
+            },
+        },
+        'desirable_weight': {
+            'label': {
+                'zh': 'KTO符合项系数',
+                'en': 'KTO desirable ratio'
+            },
+        },
+        'undesirable_weight': {
+            'label': {
+                'zh': 'KTO不符合项系数',
+                'en': 'KTO undesirable ratio'
+            },
+        }
+    }
+
+    @classmethod
+    def do_build_ui(cls, base_tab: Type['BaseUI']):
+        with gr.Accordion(elem_id='rlhf_tab', open=False):
+            with gr.Blocks():
+                with gr.Row():
+                    rlhf_type = gr.Dropdown(elem_id='rlhf_type')
+                    ref_model_type = gr.Dropdown(
+                        elem_id='ref_model_type',
+                        choices=ModelType.get_model_name_list() + cls.get_custom_name_list(),
+                        scale=20)
+                    ref_model_id_or_path = gr.Textbox(elem_id='ref_model_id_or_path', lines=1, scale=20)
+                    model_state = gr.State({})
+                with gr.Row():
+                    loss_type = gr.Dropdown(elem_id='loss_type')
+                    gr.Textbox(elem_id='max_prompt_length', lines=1, scale=20)
+                    beta = gr.Slider(elem_id='beta', minimum=0., maximum=5.0, step=0.1, scale=20)
+                    gr.Slider(elem_id='sft_beta', minimum=0., maximum=0.95, step=0.05, scale=20)
+                    gr.Slider(elem_id='simpo_gamma', minimum=0., maximum=2.0, step=0.1, scale=20)
+                    gr.Slider(elem_id='desirable_weight', minimum=0., maximum=2.0, step=0.1, scale=20)
+                    gr.Slider(elem_id='undesirable_weight', minimum=0., maximum=2.0, step=0.1, scale=20)
+
+            def update_input_model(choice, model_state=None):
+                if choice is None:
+                    return None
+                if model_state and choice in model_state:
+                    model_id_or_path = model_state[choice]
+                else:
+                    model_id_or_path = MODEL_MAPPING[choice]['model_id_or_path']
+                return model_id_or_path
+
+            def update_model_id_or_path(model_type, model_id_or_path, model_state):
+                if model_type is None or isinstance(model_type, list):
+                    return model_state
+                model_state[model_type] = model_id_or_path
+                return model_state
+
+            def update_value(rlhf_type):
+                beta = None
+                if rlhf_type in ['dpo', 'orpo', 'kto', 'cpo']:
+                    beta = 0.1
+                elif rlhf_type == 'simpo':
+                    beta = 2.0
+
+                loss_type = None
+                if rlhf_type in ['dpo', 'cpo']:
+                    loss_type = 'sigmoid'
+                elif rlhf_type == 'kto':
+                    loss_type = 'kto'
+
+                return beta, loss_type
+
+            rlhf_type.change(update_value, inputs=[rlhf_type], outputs=[beta, loss_type])
+
+            ref_model_type.change(
+                update_input_model, inputs=[ref_model_type, model_state], outputs=[ref_model_id_or_path])
+
+            ref_model_id_or_path.change(
+                update_model_id_or_path,
+                inputs=[ref_model_type, ref_model_id_or_path, model_state],
+                outputs=[model_state])
diff --git a/swift/ui/llm_train/runtime.py b/swift/ui/llm_train/runtime.py
