Merge branch 'main' into dev-2

llama_lora/lib/finetune.py

@@ -57,6 +57,9 @@ def train(
     save_steps: int = 200,
     save_total_limit: int = 3,
     logging_steps: int = 10,
+    #
+    additional_training_arguments: Union[dict, str, None] = None,
+    additional_lora_config: Union[dict, str, None] = None,
     # logging
     callbacks: List[Any] = [],
     # wandb params
@@ -70,6 +73,27 @@ def train(
 ):
     if lora_modules_to_save is not None and len(lora_modules_to_save) <= 0:
         lora_modules_to_save = None
+
+    if isinstance(additional_training_arguments, str):
+        additional_training_arguments = additional_training_arguments.strip()
+    if not additional_training_arguments:
+        additional_training_arguments = None
+    if isinstance(additional_training_arguments, str):
+        try:
+            additional_training_arguments = json.loads(additional_training_arguments)
+        except Exception as e:
+            raise ValueError(f"Could not parse additional_training_arguments: {e}")
+
+    if isinstance(additional_lora_config, str):
+        additional_lora_config = additional_lora_config.strip()
+    if not additional_lora_config:
+        additional_lora_config = None
+    if isinstance(additional_lora_config, str):
+        try:
+            additional_lora_config = json.loads(additional_lora_config)
+        except Exception as e:
+            raise ValueError(f"Could not parse additional_training_arguments: {e}")
+
     # for logging
     finetune_args = {
         'micro_batch_size': micro_batch_size,
@@ -92,6 +116,8 @@ def train(
         'save_steps': save_steps,
         'save_total_limit': save_total_limit,
         'logging_steps': logging_steps,
+        'additional_training_arguments': additional_training_arguments,
+        'additional_lora_config': additional_lora_config,
     }
     if val_set_size and val_set_size > 0:
         finetune_args['val_set_size'] = val_set_size
@@ -244,6 +270,7 @@ def train(
         lora_dropout=lora_dropout,
         bias="none",
         task_type="CAUSAL_LM",
+        **additional_lora_config,
     )
     model = get_peft_model(model, config)
     if bf16:
@@ -337,6 +364,7 @@ def train(
             group_by_length=group_by_length,
             report_to="wandb" if use_wandb else None,
             run_name=wandb_run_name if use_wandb else None,
+            **additional_training_arguments
         ),
         data_collator=transformers.DataCollatorForSeq2Seq(
             tokenizer, pad_to_multiple_of=8, return_tensors="pt", padding=True

llama_lora/ui/finetune_ui.py

@@ -305,6 +305,8 @@ def do_train(
     save_steps,
     save_total_limit,
     logging_steps,
+    additional_training_arguments,
+    additional_lora_config,
     model_name,
     continue_from_model,
     continue_from_checkpoint,
@@ -566,6 +568,8 @@ Train data (first 10):
             save_steps=save_steps,
             save_total_limit=save_total_limit,
             logging_steps=logging_steps,
+            additional_training_arguments=additional_training_arguments,
+            additional_lora_config=additional_lora_config,
             callbacks=training_callbacks,
             wandb_api_key=Global.wandb_api_key,
             wandb_project=Global.default_wandb_project if Global.enable_wandb else None,
@@ -632,6 +636,8 @@ def handle_load_params_from_model(
     save_steps,
     save_total_limit,
     logging_steps,
+    additional_training_arguments,
+    additional_lora_config,
     lora_target_module_choices,
     lora_modules_to_save_choices,
 ):
@@ -706,6 +712,16 @@ def handle_load_params_from_model(
                 save_total_limit = value
             elif key == "logging_steps":
                 logging_steps = value
+            elif key == "additional_training_arguments":
+                if value:
+                    additional_training_arguments = json.dumps(value, indent=2)
+                else:
+                    additional_training_arguments = ""
+            elif key == "additional_lora_config":
+                if value:
+                    additional_lora_config = json.dumps(value, indent=2)
+                else:
+                    additional_lora_config = ""
             elif key == "group_by_length":
                 pass
             elif key == "resume_from_checkpoint":
@@ -748,6 +764,8 @@ def handle_load_params_from_model(
         save_steps,
         save_total_limit,
         logging_steps,
+        additional_training_arguments,
+        additional_lora_config,
         lora_target_module_choices,
         lora_modules_to_save_choices
     )
@@ -946,13 +964,14 @@ def finetune_ui():
                     info="The initial learning rate for the optimizer. A higher learning rate may speed up convergence but also cause instability or divergence. A lower learning rate may require more steps to reach optimal performance but also avoid overshooting or oscillating around local minima."
                 )
 
-                with gr.Column():
+                with gr.Column(elem_id="finetune_eval_data_group"):
                     evaluate_data_count = gr.Slider(
                         minimum=0, maximum=1, step=1, value=0,
                         label="Evaluation Data Count",
                         info="The number of data to be used for evaluation. This specific amount of data will be randomly chosen from the training dataset for evaluating the model's performance during the process, without contributing to the actual training.",
                         elem_id="finetune_evaluate_data_count"
                     )
+                gr.HTML(elem_classes="flex_vertical_grow_area")
 
                 with gr.Box(elem_id="finetune_continue_from_model_box"):
                     with gr.Row():
@@ -996,6 +1015,18 @@ def finetune_ui():
                         bf16 = gr.Checkbox(label="BF16", value=False)
                         gradient_checkpointing = gr.Checkbox(
                             label="gradient_checkpointing", value=False)
+                    with gr.Column(variant="panel", elem_id="finetune_additional_training_arguments_box"):
+                        gr.Textbox(
+                            label="Additional Training Arguments",
+                            info="Additional training arguments to be passed to the Trainer in JSON format. Note that this can override ALL other arguments set elsewhere. See https://bit.ly/hf20-transformers-training-arguments for more details.",
+                            elem_id="finetune_additional_training_arguments_textbox_for_label_display"
+                            )
+                        additional_training_arguments = gr.Code(
+                            show_label=False,
+                            language="json",
+                            value="",
+                            # lines=2,
+                            elem_id="finetune_additional_training_arguments")
 
             with gr.Column():
                 lora_r = gr.Slider(
@@ -1077,8 +1108,20 @@ def finetune_ui():
                                      lora_modules_to_save_add, lora_modules_to_save],
                         ))
 
-                # with gr.Column():
-                #     pass
+                    with gr.Column(variant="panel", elem_id="finetune_additional_lora_config_box"):
+                        gr.Textbox(
+                            label="Additional LoRA Config",
+                            info="Additional LoraConfig in JSON format. Note that this can override ALL other arguments set elsewhere.",
+                            elem_id="finetune_additional_lora_config_textbox_for_label_display"
+                            )
+                        additional_lora_config = gr.Code(
+                            show_label=False,
+                            language="json",
+                            value="",
+                            # lines=2,
+                            elem_id="finetune_additional_lora_config")
+
+                gr.HTML(elem_classes="flex_vertical_grow_area no_limit")
 
                 with gr.Column(elem_id="finetune_log_and_save_options_group_container"):
                     with gr.Row(elem_id="finetune_log_and_save_options_group"):
@@ -1177,6 +1220,8 @@ def finetune_ui():
             save_steps,
             save_total_limit,
             logging_steps,
+            additional_training_arguments,
+            additional_lora_config,
         ]
 
         things_that_might_timeout.append(

llama_lora/ui/main_page.py

@@ -250,6 +250,15 @@ def main_page_custom_css():
         display: none;
     }
 
+    .flex_vertical_grow_area {
+        margin-top: calc(var(--layout-gap) * -1) !important;
+        flex-grow: 1 !important;
+        max-height: calc(var(--layout-gap) * 2);
+    }
+    .flex_vertical_grow_area.no_limit {
+        max-height: unset;
+    }
+
     #page_title {
         flex-grow: 3;
     }
@@ -808,13 +817,32 @@ def main_page_custom_css():
     }
 
     #finetune_log_and_save_options_group_container {
-        flex-grow: 1 !important;
-        justify-content: flex-end;
+        flex-grow: 0 !important;
     }
     #finetune_model_name_group {
         flex-grow: 0 !important;
     }
 
+    #finetune_eval_data_group {
+        flex-grow: 0 !important;
+    }
+
+    #finetune_additional_training_arguments_box > .form,
+    #finetune_additional_lora_config_box > .form {
+        border: 0;
+        background: transparent;
+    }
+    #finetune_additional_training_arguments_textbox_for_label_display,
+    #finetune_additional_lora_config_textbox_for_label_display {
+        padding: 0;
+        margin-bottom: -10px;
+        background: transparent;
+    }
+    #finetune_additional_training_arguments_textbox_for_label_display textarea,
+    #finetune_additional_lora_config_textbox_for_label_display textarea {
+        display: none;
+    }
+
     @media screen and (max-width: 392px) {
         #inference_lora_model, #inference_lora_model_group, #finetune_template {
             border-bottom-left-radius: 0;