add more finetune options

llama_lora/lib/finetune.py

@@ -28,8 +28,11 @@ def train(
     tokenizer: Any,
     output_dir: str,
     train_data: List[Any],
+    #
     load_in_8bit=True,
     fp16=True,
+    bf16=False,
+    gradient_checkpointing=False,
     # training hyperparams
     micro_batch_size: int = 4,
     gradient_accumulation_steps: int = 32,
@@ -79,18 +82,21 @@ def train(
         'lora_alpha': lora_alpha,
         'lora_dropout': lora_dropout,
         'lora_target_modules': lora_target_modules,
+        'lora_modules_to_save': lora_modules_to_save or [],
         'train_on_inputs': train_on_inputs,
         'group_by_length': group_by_length,
         'load_in_8bit': load_in_8bit,
         'fp16': fp16,
+        'bf16': bf16,
+        'gradient_checkpointing': gradient_checkpointing,
         'save_steps': save_steps,
         'save_total_limit': save_total_limit,
         'logging_steps': logging_steps,
     }
     if val_set_size and val_set_size > 0:
         finetune_args['val_set_size'] = val_set_size
-    if lora_modules_to_save:
-        finetune_args['lora_modules_to_save'] = lora_modules_to_save
+    # if lora_modules_to_save:
+    #     finetune_args['lora_modules_to_save'] = lora_modules_to_save
     if resume_from_checkpoint:
         finetune_args['resume_from_checkpoint'] = resume_from_checkpoint
 
@@ -232,6 +238,8 @@ def train(
         task_type="CAUSAL_LM",
     )
     model = get_peft_model(model, config)
+    if bf16:
+        model = model.to(torch.bfloat16)
 
     # If train_data is a list, convert it to datasets.Dataset
     if isinstance(train_data, list):
@@ -289,11 +297,13 @@ def train(
         # https://huggingface.co/docs/transformers/main/en/main_classes/trainer#transformers.TrainingArguments
         args=transformers.TrainingArguments(
             per_device_train_batch_size=micro_batch_size,
+            gradient_checkpointing=gradient_checkpointing,
             gradient_accumulation_steps=gradient_accumulation_steps,
             warmup_steps=100,
             num_train_epochs=num_train_epochs,
             learning_rate=learning_rate,
             fp16=fp16,
+            bf16=bf16,
             logging_steps=logging_steps,
             optim="adamw_torch",
             evaluation_strategy="steps" if val_set_size > 0 else "no",

llama_lora/ui/finetune_ui.py

@@ -299,6 +299,8 @@ def do_train(
     lora_modules_to_save,
     load_in_8bit,
     fp16,
+    bf16,
+    gradient_checkpointing,
     save_steps,
     save_total_limit,
     logging_steps,
@@ -393,6 +395,8 @@ Train options: {json.dumps({
     'lora_modules_to_save': lora_modules_to_save,
     'load_in_8bit': load_in_8bit,
     'fp16': fp16,
+    'bf16': bf16,
+    'gradient_checkpointing': gradient_checkpointing,
     'model_name': model_name,
     'continue_from_model': continue_from_model,
     'continue_from_checkpoint': continue_from_checkpoint,
@@ -532,6 +536,8 @@ Train data (first 10):
             train_on_inputs=train_on_inputs,
             load_in_8bit=load_in_8bit,
             fp16=fp16,
+            bf16=bf16,
+            gradient_checkpointing=gradient_checkpointing,
             group_by_length=False,
             resume_from_checkpoint=resume_from_checkpoint,
             save_steps=save_steps,
@@ -548,8 +554,9 @@ Train data (first 10):
         logs_str = "\n".join([json.dumps(log)
                              for log in log_history]) or "None"
 
-        result_message = f"Training ended:\n{str(train_output)}\n\nLogs:\n{logs_str}"
+        result_message = f"Training ended:\n{str(train_output)}"
         print(result_message)
+        # result_message += f"\n\nLogs:\n{logs_str}"
 
         clear_cache()
 
@@ -597,6 +604,8 @@ def handle_load_params_from_model(
     lora_modules_to_save,
     load_in_8bit,
     fp16,
+    bf16,
+    gradient_checkpointing,
     save_steps,
     save_total_limit,
     logging_steps,
@@ -650,18 +659,24 @@ def handle_load_params_from_model(
                 lora_dropout = value
             elif key == "lora_target_modules":
                 lora_target_modules = value
-                for element in value:
-                    if element not in lora_target_module_choices:
-                        lora_target_module_choices.append(element)
+                if value:
+                    for element in value:
+                        if element not in lora_target_module_choices:
+                            lora_target_module_choices.append(element)
             elif key == "lora_modules_to_save":
                 lora_modules_to_save = value
-                for element in value:
-                    if element not in lora_modules_to_save_choices:
-                        lora_modules_to_save_choices.append(element)
+                if value:
+                    for element in value:
+                        if element not in lora_modules_to_save_choices:
+                            lora_modules_to_save_choices.append(element)
             elif key == "load_in_8bit":
                 load_in_8bit = value
             elif key == "fp16":
                 fp16 = value
+            elif key == "bf16":
+                bf16 = value
+            elif key == "gradient_checkpointing":
+                gradient_checkpointing = value
             elif key == "save_steps":
                 save_steps = value
             elif key == "save_total_limit":
@@ -705,6 +720,8 @@ def handle_load_params_from_model(
             value=lora_modules_to_save, choices=lora_modules_to_save_choices),
         load_in_8bit,
         fp16,
+        bf16,
+        gradient_checkpointing,
         save_steps,
         save_total_limit,
         logging_steps,
@@ -949,9 +966,11 @@ def finetune_ui():
                     )
 
                 with gr.Accordion("Advanced Options", open=False, elem_id="finetune_advance_options_accordion"):
-                    with gr.Row():
-                        load_in_8bit = gr.Checkbox(label="8bit", value=True)
+                    with gr.Row(elem_id="finetune_advanced_options_checkboxes"):
+                        load_in_8bit = gr.Checkbox(label="8bit", value=False)
                         fp16 = gr.Checkbox(label="FP16", value=True)
+                        bf16 = gr.Checkbox(label="BF16", value=False)
+                        gradient_checkpointing = gr.Checkbox(label="gradient_checkpointing", value=False)
 
             with gr.Column():
                 lora_r = gr.Slider(
@@ -1002,57 +1021,62 @@ def finetune_ui():
                                  lora_target_modules_add, lora_target_modules],
                     ))
 
-                with gr.Column(elem_id="finetune_lora_modules_to_save_box"):
-                    lora_modules_to_save = gr.CheckboxGroup(
-                        label="LoRA Modules To Save",
-                        choices=default_lora_modules_to_save_choices,
-                        value=[],
-                        # info="",
-                        elem_id="finetune_lora_modules_to_save"
-                    )
-                    lora_modules_to_save_choices = gr.State(
-                        value=default_lora_modules_to_save_choices)
-                    with gr.Box(elem_id="finetune_lora_modules_to_save_add_box"):
-                        with gr.Row():
-                            lora_modules_to_save_add = gr.Textbox(
-                                lines=1, max_lines=1, show_label=False,
-                                elem_id="finetune_lora_modules_to_save_add"
-                            )
-                            lora_modules_to_save_add_btn = gr.Button(
-                                "Add",
-                                elem_id="finetune_lora_modules_to_save_add_btn"
-                            )
-                            lora_modules_to_save_add_btn.style(
-                                full_width=False, size="sm")
-                    things_that_might_timeout.append(lora_modules_to_save_add_btn.click(
-                        handle_lora_modules_to_save_add,
-                        inputs=[lora_modules_to_save_choices,
-                                lora_modules_to_save_add, lora_modules_to_save],
-                        outputs=[lora_modules_to_save_choices,
-                                 lora_modules_to_save_add, lora_modules_to_save],
-                    ))
-
-                with gr.Row():
-                    logging_steps = gr.Number(
-                        label="Logging Steps",
-                        precision=0,
-                        value=10,
-                        elem_id="finetune_logging_steps"
-                    )
-                    save_steps = gr.Number(
-                        label="Steps Per Save",
-                        precision=0,
-                        value=500,
-                        elem_id="finetune_save_steps"
-                    )
-                    save_total_limit = gr.Number(
-                        label="Saved Checkpoints Limit",
-                        precision=0,
-                        value=5,
-                        elem_id="finetune_save_total_limit"
-                    )
+                with gr.Accordion("Advanced LoRA Options", open=False, elem_id="finetune_advance_lora_options_accordion"):
+                    with gr.Column(elem_id="finetune_lora_modules_to_save_box"):
+                        lora_modules_to_save = gr.CheckboxGroup(
+                            label="LoRA Modules To Save",
+                            choices=default_lora_modules_to_save_choices,
+                            value=[],
+                            # info="",
+                            elem_id="finetune_lora_modules_to_save"
+                        )
+                        lora_modules_to_save_choices = gr.State(
+                            value=default_lora_modules_to_save_choices)
+                        with gr.Box(elem_id="finetune_lora_modules_to_save_add_box"):
+                            with gr.Row():
+                                lora_modules_to_save_add = gr.Textbox(
+                                    lines=1, max_lines=1, show_label=False,
+                                    elem_id="finetune_lora_modules_to_save_add"
+                                )
+                                lora_modules_to_save_add_btn = gr.Button(
+                                    "Add",
+                                    elem_id="finetune_lora_modules_to_save_add_btn"
+                                )
+                                lora_modules_to_save_add_btn.style(
+                                    full_width=False, size="sm")
+                        things_that_might_timeout.append(lora_modules_to_save_add_btn.click(
+                            handle_lora_modules_to_save_add,
+                            inputs=[lora_modules_to_save_choices,
+                                    lora_modules_to_save_add, lora_modules_to_save],
+                            outputs=[lora_modules_to_save_choices,
+                                     lora_modules_to_save_add, lora_modules_to_save],
+                        ))
+
+                # with gr.Column():
+                #     pass
+
+                with gr.Column(elem_id="finetune_log_and_save_options_group_container"):
+                    with gr.Row(elem_id="finetune_log_and_save_options_group"):
+                        logging_steps = gr.Number(
+                            label="Logging Steps",
+                            precision=0,
+                            value=10,
+                            elem_id="finetune_logging_steps"
+                        )
+                        save_steps = gr.Number(
+                            label="Steps Per Save",
+                            precision=0,
+                            value=500,
+                            elem_id="finetune_save_steps"
+                        )
+                        save_total_limit = gr.Number(
+                            label="Saved Checkpoints Limit",
+                            precision=0,
+                            value=5,
+                            elem_id="finetune_save_total_limit"
+                        )
 
-                with gr.Column():
+                with gr.Column(elem_id="finetune_model_name_group"):
                     model_name = gr.Textbox(
                         lines=1, label="LoRA Model Name", value=random_name,
                         max_lines=1,
@@ -1123,6 +1147,8 @@ def finetune_ui():
             lora_modules_to_save,
             load_in_8bit,
             fp16,
+            bf16,
+            gradient_checkpointing,
             save_steps,
             save_total_limit,
             logging_steps,

llama_lora/ui/main_page.py

@@ -734,11 +734,12 @@ def main_page_custom_css():
     }
 
     #finetune_lora_target_modules_box,
-    #finetune_lora_modules_to_save_box {
-        margin-top: calc((var(--layout-gap) + 8px) * -1)
+    #finetune_lora_target_modules_box + #finetune_lora_modules_to_save_box {
+        margin-top: calc((var(--layout-gap) + 8px) * -1);
+        flex-grow: 0 !important;
     }
     #finetune_lora_target_modules_box > .form,
-    #finetune_lora_modules_to_save_box > .form {
+    #finetune_lora_target_modules_box + #finetune_lora_modules_to_save_box > .form {
         padding-top: calc((var(--layout-gap) + 8px) / 3);
         border-top: 0;
         border-top-left-radius: 0;
@@ -747,7 +748,7 @@ def main_page_custom_css():
         position: relative;
     }
     #finetune_lora_target_modules_box > .form::before,
-    #finetune_lora_modules_to_save_box > .form::before {
+    #finetune_lora_target_modules_box + #finetune_lora_modules_to_save_box > .form::before {
         content: "";
         display: block;
         position: absolute;
@@ -802,6 +803,18 @@ def main_page_custom_css():
         padding: 4px 8px;
     }
 
+    #finetune_advanced_options_checkboxes > * > * {
+        min-width: auto;
+    }
+
+    #finetune_log_and_save_options_group_container {
+        flex-grow: 1 !important;
+        justify-content: flex-end;
+    }
+    #finetune_model_name_group {
+        flex-grow: 0 !important;
+    }
+
     @media screen and (max-width: 392px) {
         #inference_lora_model, #inference_lora_model_group, #finetune_template {
             border-bottom-left-radius: 0;