make the training process async

README.md

@@ -70,7 +70,7 @@ setup: |
 # Start the app.
 run: |
   echo 'Starting...'
-  python llama_lora_tuner/app.py --data_dir='/data' --wandb_api_key="$([ -f /data/secrets/wandb_api_key ] && cat /data/secrets/wandb_api_key | tr -d '\n')" --base_model=decapoda-research/llama-7b-hf --base_model_choices='decapoda-research/llama-7b-hf,nomic-ai/gpt4all-j,databricks/dolly-v2-7b --share
+  python llama_lora_tuner/app.py --data_dir='/data' --wandb_api_key="$([ -f /data/secrets/wandb_api_key ] && cat /data/secrets/wandb_api_key | tr -d '\n')" --timezone='Atlantic/Reykjavik' --base_model=decapoda-research/llama-7b-hf --base_model_choices='decapoda-research/llama-7b-hf,nomic-ai/gpt4all-j,databricks/dolly-v2-7b --share
 ```
 
 Then launch a cluster to run the task:
@@ -100,7 +100,7 @@ When you are done, run `sky stop <cluster_name>` to stop the cluster. To termina
 
 ```bash
 pip install -r requirements.lock.txt
-python app.py --data_dir='./data' --base_model='decapoda-research/llama-7b-hf' --share
+python app.py --data_dir='./data' --base_model='decapoda-research/llama-7b-hf' --timezone='Atlantic/Reykjavik' --share
 ```
 
 You will see the local and public URLs of the app in the terminal. Open the URL in your browser to use the app.

app.py

@@ -28,6 +28,7 @@ def main(
     ui_dev_mode: Union[bool, None] = None,
     wandb_api_key: Union[str, None] = None,
     wandb_project: Union[str, None] = None,
+    timezone: Union[str, None] = None,
 ):
     '''
     Start the LLaMA-LoRA Tuner UI.
@@ -76,6 +77,9 @@ def main(
     if wandb_project is not None:
         Config.default_wandb_project = wandb_project
 
+    if timezone is not None:
+        Config.timezone = timezone
+
     if ui_dev_mode is not None:
         Config.ui_dev_mode = ui_dev_mode
 

config.yaml.sample

@@ -9,6 +9,8 @@ base_model_choices:
 load_8bit: false
 trust_remote_code: false
 
+# timezone: Atlantic/Reykjavik
+
 # UI Customization
 # ui_title: LLM Tuner
 # ui_emoji: 🦙🎛️

llama_lora/config.py

@@ -1,5 +1,6 @@
 import os
-from typing import List, Union
+import pytz
+from typing import List, Union, Any
 
 
 class Config:
@@ -15,6 +16,8 @@ class Config:
 
     trust_remote_code: bool = False
 
+    timezone: Any = pytz.UTC
+
     # WandB
     enable_wandb: Union[bool, None] = False
     wandb_api_key: Union[str, None] = None
@@ -37,6 +40,9 @@ def process_config():
         base_model_choices = [name.strip() for name in base_model_choices]
         Config.base_model_choices = base_model_choices
 
+    if isinstance(Config.timezone, str):
+        Config.timezone = pytz.timezone(Config.timezone)
+
     if Config.default_base_model_name not in Config.base_model_choices:
         Config.base_model_choices = [Config.default_base_model_name] + Config.base_model_choices
 

llama_lora/globals.py

@@ -12,6 +12,7 @@ import nvidia_smi
 from .dynamic_import import dynamic_import
 from .config import Config
 from .utils.lru_cache import LRUCache
+from .utils.eta_predictor import ETAPredictor
 
 
 class Global:
@@ -31,6 +32,24 @@ class Global:
     # Training Control
     should_stop_training: bool = False
 
+    # Training Status
+    is_train_starting: bool = False
+    is_training: bool = False
+    train_started_at: float = 0.0
+    training_error_message: Union[str, None] = None
+    training_error_detail: Union[str, None] = None
+    training_total_epochs: int = 0
+    training_current_epoch: float = 0.0
+    training_total_steps: int = 0
+    training_current_step: int = 0
+    training_progress: float = 0.0
+    training_log_history: List[Any] = []
+    training_status_text: str = ""
+    training_eta_predictor = ETAPredictor()
+    training_eta: Union[int, None] = None
+    train_output: Union[None, Any] = None
+    train_output_str: Union[None, str] = None
+
     # Generation Control
     should_stop_generating: bool = False
     generation_force_stopped_at: Union[float, None] = None

llama_lora/models.py

@@ -26,6 +26,8 @@ def get_peft_model_class():
 def get_new_base_model(base_model_name):
     if Config.ui_dev_mode:
         return
+    if Global.is_train_starting or Global.is_training:
+        raise Exception("Cannot load new base model while training.")
 
     if Global.new_base_model_that_is_ready_to_be_used:
         if Global.name_of_new_base_model_that_is_ready_to_be_used == base_model_name:
@@ -121,6 +123,9 @@ def get_tokenizer(base_model_name):
     if Config.ui_dev_mode:
         return
 
+    if Global.is_train_starting or Global.is_training:
+        raise Exception("Cannot load new base model while training.")
+
     loaded_tokenizer = Global.loaded_tokenizers.get(base_model_name)
     if loaded_tokenizer:
         return loaded_tokenizer
@@ -150,6 +155,9 @@ def get_model(
     if Config.ui_dev_mode:
         return
 
+    if Global.is_train_starting or Global.is_training:
+        raise Exception("Cannot load new base model while training.")
+
     if peft_model_name == "None":
         peft_model_name = None
 

llama_lora/ui/finetune/finetune_ui.py

@@ -27,7 +27,8 @@ from .previewing import (
     refresh_dataset_items_count,
 )
 from .training import (
-    do_train
+    do_train,
+    render_training_status
 )
 
 register_css_style('finetune', relative_read_file(__file__, "style.css"))
@@ -770,19 +771,22 @@ def finetune_ui():
             )
         )
 
-        train_output = gr.Text(
+        train_status = gr.HTML(
             "Training results will be shown here.",
             label="Train Output",
             elem_id="finetune_training_status")
 
-        train_progress = train_btn.click(
+        training_indicator = gr.HTML(
+            "training_indicator", visible=False, elem_id="finetune_training_indicator")
+
+        train_start = train_btn.click(
             fn=do_train,
             inputs=(dataset_inputs + finetune_args + [
                 model_name,
                 continue_from_model,
                 continue_from_checkpoint,
             ]),
-            outputs=train_output
+            outputs=[train_status, training_indicator]
         )
 
         # controlled by JS, shows the confirm_abort_button
@@ -790,13 +794,20 @@ def finetune_ui():
         confirm_abort_button.click(
             fn=do_abort_training,
             inputs=None, outputs=None,
-            cancels=[train_progress])
-
-        stop_timeoutable_btn = gr.Button(
-            "stop not-responding elements",
-            elem_id="inference_stop_timeoutable_btn",
-            elem_classes="foot_stop_timeoutable_btn")
-        stop_timeoutable_btn.click(
-            fn=None, inputs=None, outputs=None, cancels=things_that_might_timeout)
+            cancels=[train_start])
 
+    training_status_updates = finetune_ui_blocks.load(
+        fn=render_training_status,
+        inputs=None,
+        outputs=[train_status, training_indicator],
+        every=0.1
+    )
     finetune_ui_blocks.load(_js=relative_read_file(__file__, "script.js"))
+
+    # things_that_might_timeout.append(training_status_updates)
+    stop_timeoutable_btn = gr.Button(
+        "stop not-responding elements",
+        elem_id="inference_stop_timeoutable_btn",
+        elem_classes="foot_stop_timeoutable_btn")
+    stop_timeoutable_btn.click(
+        fn=None, inputs=None, outputs=None, cancels=things_that_might_timeout)

llama_lora/ui/finetune/script.js

@@ -130,10 +130,10 @@ function finetune_ui_blocks_js() {
 
   // Show/hide start and stop button base on the state.
   setTimeout(function () {
-    // Make the '#finetune_training_status > .wrap' element appear
-    if (!document.querySelector('#finetune_training_status > .wrap')) {
-      document.getElementById('finetune_confirm_stop_btn').click();
-    }
+    // Make the '#finetune_training_indicator > .wrap' element appear
+    // if (!document.querySelector('#finetune_training_indicator > .wrap')) {
+    //   document.getElementById('finetune_confirm_stop_btn').click();
+    // }
 
     setTimeout(function () {
       let resetStopButtonTimer;
@@ -156,11 +156,20 @@ function finetune_ui_blocks_js() {
           document.getElementById('finetune_confirm_stop_btn').style.display =
             'block';
         });
-      const output_wrap_element = document.querySelector(
-        '#finetune_training_status > .wrap'
+      // const training_indicator_wrap_element = document.querySelector(
+      //   '#finetune_training_indicator > .wrap'
+      // );
+      const training_indicator_element = document.querySelector(
+        '#finetune_training_indicator'
       );
-      function handle_output_wrap_element_class_change() {
-        if (Array.from(output_wrap_element.classList).includes('hide')) {
+      let isTraining = undefined;
+      function handle_training_indicator_change() {
+        // const wrapperHidden = Array.from(training_indicator_wrap_element.classList).includes('hide');
+        const hidden = Array.from(training_indicator_element.classList).includes('hidden');
+        const newIsTraining = !(/* wrapperHidden && */ hidden);
+        if (newIsTraining === isTraining) return;
+        isTraining = newIsTraining;
+        if (!isTraining) {
           if (resetStopButtonTimer) clearTimeout(resetStopButtonTimer);
           document.getElementById('finetune_start_btn').style.display = 'block';
           document.getElementById('finetune_stop_btn').style.display = 'none';
@@ -173,13 +182,19 @@ function finetune_ui_blocks_js() {
             'none';
         }
       }
+      // new MutationObserver(function (mutationsList, observer) {
+      //   handle_training_indicator_change();
+      // }).observe(training_indicator_wrap_element, {
+      //   attributes: true,
+      //   attributeFilter: ['class'],
+      // });
       new MutationObserver(function (mutationsList, observer) {
-        handle_output_wrap_element_class_change();
-      }).observe(output_wrap_element, {
+        handle_training_indicator_change();
+      }).observe(training_indicator_element, {
         attributes: true,
         attributeFilter: ['class'],
       });
-      handle_output_wrap_element_class_change();
+      handle_training_indicator_change();
     }, 500);
   }, 0);
 }

llama_lora/ui/finetune/style.css

@@ -255,8 +255,118 @@
     display: none;
 }
 
-/* in case if there's too many logs on the previous run and made the box too high */
-#finetune_training_status:has(.wrap:not(.hide)) {
-    max-height: 160px;
-    height: 160px;
+#finetune_training_status > .wrap {
+    border: 0;
+    background: transparent;
+    pointer-events: none;
+    top: 0;
+    bottom: 0;
+    left: 0;
+    right: 0;
+}
+#finetune_training_status > .wrap .meta-text-center {
+    transform: none !important;
+}
+
+#finetune_training_status .progress-block {
+    min-height: 100px;
+    display: flex;
+    justify-content: center;
+    align-items: center;
+    background: var(--panel-background-fill);
+    border-radius: var(--radius-lg);
+    border: var(--block-border-width) solid var(--border-color-primary);
+    padding: var(--block-padding);
+}
+#finetune_training_status .progress-block.is_training {
+    min-height: 160px;
+}
+#finetune_training_status .progress-block .empty-text {
+    text-transform: uppercase;
+    font-weight: 700;
+    font-size: 120%;
+    opacity: 0.12;
+}
+#finetune_training_status .progress-block .meta-text {
+    position: absolute;
+    top: 0;
+    right: 0;
+    z-index: var(--layer-2);
+    padding: var(--size-1) var(--size-2);
+    font-size: var(--text-sm);
+    font-family: var(--font-mono);
+}
+#finetune_training_status .progress-block .status {
+    white-space: pre-wrap;
+}
+#finetune_training_status .progress-block .progress-level {
+    display: flex;
+    flex-direction: column;
+    align-items: center;
+    z-index: var(--layer-2);
+    width: var(--size-full);
+}
+#finetune_training_status .progress-block .progress-level-inner {
+    margin: var(--size-2) auto;
+    color: var(--body-text-color);
+    font-size: var(--text-sm);
+    font-family: var(--font-mono);
+}
+#finetune_training_status .progress-block .progress-bar-wrap {
+    border: 1px solid var(--border-color-primary);
+    background: var(--background-fill-primary);
+    width: 55.5%;
+    height: var(--size-4);
+}
+#finetune_training_status .progress-block .progress-bar {
+    transform-origin: left;
+    background-color: var(--loader-color);
+    width: var(--size-full);
+    height: var(--size-full);
+    transition: all 150ms ease 0s;
 }
+
+#finetune_training_status .progress-block .output {
+    display: flex;
+    flex-direction: column;
+    justify-content: center;
+    align-items: center;
+}
+#finetune_training_status .progress-block .output .title {
+    padding: var(--size-1) var(--size-3);
+    font-weight: var(--weight-bold);
+    font-size: var(--text-lg);
+    line-height: var(--line-xs);
+}
+#finetune_training_status .progress-block .output .message {
+    padding: var(--size-1) var(--size-3);
+    color: var(--body-text-color) !important;
+    font-family: var(--font-mono);
+    white-space: pre-wrap;
+}
+
+#finetune_training_status .progress-block .error {
+    display: flex;
+    flex-direction: column;
+    justify-content: center;
+    align-items: center;
+}
+#finetune_training_status .progress-block .error .title {
+    padding: var(--size-1) var(--size-3);
+    color: var(--color-red-500);
+    font-weight: var(--weight-bold);
+    font-size: var(--text-lg);
+    line-height: var(--line-xs);
+}
+#finetune_training_status .progress-block .error .error-message {
+    padding: var(--size-1) var(--size-3);
+    color: var(--body-text-color) !important;
+    font-family: var(--font-mono);
+    white-space: pre-wrap;
+}
+#finetune_training_status .progress-block.is_error {
+    /* background: var(--error-background-fill) !important; */
+    border: 1px solid var(--error-border-color) !important;
+}
+
+#finetune_training_indicator { display: none; }

llama_lora/ui/finetune/training.py

@@ -1,24 +1,26 @@
 import os
 import json
 import time
+import datetime
+import pytz
+import socket
+import threading
+import traceback
 import gradio as gr
-import math
 
-from transformers import TrainerCallback
 from huggingface_hub import try_to_load_from_cache, snapshot_download
 
 from ...config import Config
 from ...globals import Global
 from ...models import clear_cache, unload_models
 from ...utils.prompter import Prompter
+from ..trainer_callback import (
+    UiTrainerCallback, reset_training_status,
+    update_training_states, set_train_output
+)
 
 from .data_processing import get_data_from_input
 
-should_training_progress_track_tqdm = True
-
-if Global.gpu_total_cores is not None and Global.gpu_total_cores > 2560:
-    should_training_progress_track_tqdm = False
-
 
 def do_train(
     # Dataset
@@ -55,8 +57,14 @@ def do_train(
     model_name,
     continue_from_model,
     continue_from_checkpoint,
-    progress=gr.Progress(track_tqdm=should_training_progress_track_tqdm),
+    progress=gr.Progress(track_tqdm=False),
 ):
+    if Global.is_training:
+        return render_training_status()
+
+    reset_training_status()
+    Global.is_train_starting = True
+
     try:
         base_model_name = Global.base_model_name
         tokenizer_name = Global.tokenizer_name or Global.base_model_name
@@ -115,18 +123,47 @@ def do_train(
                 raise ValueError(
                     f"The output directory already exists and is not empty. ({output_dir})")
 
-        if not should_training_progress_track_tqdm:
-            progress(0, desc="Preparing train data...")
+        wandb_group = template
+        wandb_tags = [f"template:{template}"]
+        if load_dataset_from == "Data Dir" and dataset_from_data_dir:
+            wandb_group += f"/{dataset_from_data_dir}"
+            wandb_tags.append(f"dataset:{dataset_from_data_dir}")
 
-        # Need RAM for training
-        unload_models()
-        Global.new_base_model_that_is_ready_to_be_used = None
-        Global.name_of_new_base_model_that_is_ready_to_be_used = None
-        clear_cache()
+        finetune_args = {
+            'base_model': base_model_name,
+            'tokenizer': tokenizer_name,
+            'output_dir': output_dir,
+            'micro_batch_size': micro_batch_size,
+            'gradient_accumulation_steps': gradient_accumulation_steps,
+            'num_train_epochs': epochs,
+            'learning_rate': learning_rate,
+            'cutoff_len': max_seq_length,
+            'val_set_size': evaluate_data_count,
+            'lora_r': lora_r,
+            'lora_alpha': lora_alpha,
+            'lora_dropout': lora_dropout,
+            'lora_target_modules': lora_target_modules,
+            'lora_modules_to_save': lora_modules_to_save,
+            '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_param,
+            '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,
+            'wandb_api_key': Config.wandb_api_key,
+            'wandb_project': Config.default_wandb_project if Config.enable_wandb else None,
+            'wandb_group': wandb_group,
+            'wandb_run_name': model_name,
+            'wandb_tags': wandb_tags
+        }
 
         prompter = Prompter(template)
-        # variable_names = prompter.get_variable_names()
-
         data = get_data_from_input(
             load_dataset_from=load_dataset_from,
             dataset_text=dataset_text,
@@ -138,208 +175,234 @@ def do_train(
             prompter=prompter
         )
 
-        train_data = prompter.get_train_data_from_dataset(data)
-
-        def get_progress_text(epoch, epochs, last_loss):
-            progress_detail = f"Epoch {math.ceil(epoch)}/{epochs}"
-            if last_loss is not None:
-                progress_detail += f", Loss: {last_loss:.4f}"
-            return f"Training... ({progress_detail})"
-
-        if Config.ui_dev_mode:
-            Global.should_stop_training = False
-
-            message = f"""Currently in UI dev mode, not doing the actual training.
-
-Train options: {json.dumps({
-    'max_seq_length': max_seq_length,
-    'val_set_size': evaluate_data_count,
-    'micro_batch_size': micro_batch_size,
-    'gradient_accumulation_steps': gradient_accumulation_steps,
-    'epochs': epochs,
-    'learning_rate': learning_rate,
-    'train_on_inputs': train_on_inputs,
-    'lora_r': lora_r,
-    'lora_alpha': lora_alpha,
-    'lora_dropout': lora_dropout,
-    'lora_target_modules': lora_target_modules,
-    '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,
-    'resume_from_checkpoint_param': resume_from_checkpoint_param,
-}, indent=2)}
-
-Train data (first 10):
-{json.dumps(train_data[:10], indent=2)}
-            """
-            print(message)
-
-            for i in range(300):
-                if (Global.should_stop_training):
+        def training():
+            Global.is_training = True
+
+            try:
+                # Need RAM for training
+                unload_models()
+                Global.new_base_model_that_is_ready_to_be_used = None
+                Global.name_of_new_base_model_that_is_ready_to_be_used = None
+                clear_cache()
+
+                train_data = prompter.get_train_data_from_dataset(data)
+
+                if Config.ui_dev_mode:
+                    message = "Currently in UI dev mode, not doing the actual training."
+                    message += f"\n\nArgs: {json.dumps(finetune_args, indent=2)}"
+                    message += f"\n\nTrain data (first 5):\n{json.dumps(train_data[:5], indent=2)}"
+
+                    print(message)
+
+                    total_steps = 300
+                    for i in range(300):
+                        if (Global.should_stop_training):
+                            break
+
+                        current_step = i + 1
+                        total_epochs = 3
+                        current_epoch = i / 100
+                        log_history = []
+
+                        if (i > 20):
+                            loss = 3 + (i - 0) * (0.5 - 3) / (300 - 0)
+                            log_history = [{'loss': loss}]
+
+                        update_training_states(
+                            total_steps=total_steps,
+                            current_step=current_step,
+                            total_epochs=total_epochs,
+                            current_epoch=current_epoch,
+                            log_history=log_history
+                        )
+                        time.sleep(0.1)
+
+                    result_message = set_train_output(message)
+                    print(result_message)
+                    time.sleep(1)
+                    Global.is_training = False
                     return
-                epochs = 3
-                epoch = i / 100
-                last_loss = None
-                if (i > 20):
-                    last_loss = 3 + (i - 0) * (0.5 - 3) / (300 - 0)
-
-                progress(
-                    (i, 300),
-                    desc="(Simulate) " +
-                    get_progress_text(epoch, epochs, last_loss)
-                )
-
-                time.sleep(0.1)
-
-            time.sleep(2)
-            return message
-
-        if not should_training_progress_track_tqdm:
-            progress(
-                0, desc=f"Preparing model {base_model_name} for training...")
-
-        log_history = []
-
-        class UiTrainerCallback(TrainerCallback):
-            def _on_progress(self, args, state, control):
-                nonlocal log_history
 
-                if Global.should_stop_training:
-                    control.should_training_stop = True
-                total_steps = (
-                    state.max_steps if state.max_steps is not None else state.num_train_epochs * state.steps_per_epoch)
-                log_history = state.log_history
-                last_history = None
-                last_loss = None
-                if len(log_history) > 0:
-                    last_history = log_history[-1]
-                    last_loss = last_history.get('loss', None)
-
-                progress_detail = f"Epoch {math.ceil(state.epoch)}/{epochs}"
-                if last_loss is not None:
-                    progress_detail += f", Loss: {last_loss:.4f}"
-
-                progress(
-                    (state.global_step, total_steps),
-                    desc=f"Training... ({progress_detail})"
+                training_callbacks = [UiTrainerCallback]
+
+                if not os.path.exists(output_dir):
+                    os.makedirs(output_dir)
+
+                with open(os.path.join(output_dir, "info.json"), 'w') as info_json_file:
+                    dataset_name = "N/A (from text input)"
+                    if load_dataset_from == "Data Dir":
+                        dataset_name = dataset_from_data_dir
+
+                    info = {
+                        'base_model': base_model_name,
+                        'prompt_template': template,
+                        'dataset_name': dataset_name,
+                        'dataset_rows': len(train_data),
+                        'trained_on_machine': socket.gethostname(),
+                        'timestamp': time.time(),
+                    }
+                    if continue_from_model:
+                        info['continued_from_model'] = continue_from_model
+                        if continue_from_checkpoint:
+                            info['continued_from_checkpoint'] = continue_from_checkpoint
+
+                    if Global.version:
+                        info['tuner_version'] = Global.version
+
+                    json.dump(info, info_json_file, indent=2)
+
+                train_output = Global.finetune_train_fn(
+                    train_data=train_data,
+                    callbacks=training_callbacks,
+                    **finetune_args,
                 )
 
-            def on_epoch_begin(self, args, state, control, **kwargs):
-                self._on_progress(args, state, control)
-
-            def on_step_end(self, args, state, control, **kwargs):
-                self._on_progress(args, state, control)
-
-        training_callbacks = [UiTrainerCallback]
-
-        Global.should_stop_training = False
-
-        # Do not let other tqdm iterations interfere the progress reporting after training starts.
-        # progress.track_tqdm = False  # setting this dynamically is not working, determining if track_tqdm should be enabled based on GPU cores at start instead.
+                result_message = set_train_output(train_output)
+                print(result_message + "\n" + str(train_output))
 
-        if not os.path.exists(output_dir):
-            os.makedirs(output_dir)
+                clear_cache()
 
-        with open(os.path.join(output_dir, "info.json"), 'w') as info_json_file:
-            dataset_name = "N/A (from text input)"
-            if load_dataset_from == "Data Dir":
-                dataset_name = dataset_from_data_dir
+                Global.is_training = False
 
-            info = {
-                'base_model': base_model_name,
-                'prompt_template': template,
-                'dataset_name': dataset_name,
-                'dataset_rows': len(train_data),
-                'timestamp': time.time(),
+            except Exception as e:
+                traceback.print_exc()
+                Global.training_error_message = str(e)
+            finally:
+                Global.is_training = False
 
-                # These will be saved in another JSON file by the train function
-                # 'max_seq_length': max_seq_length,
-                # 'train_on_inputs': train_on_inputs,
-
-                # 'micro_batch_size': micro_batch_size,
-                # 'gradient_accumulation_steps': gradient_accumulation_steps,
-                # 'epochs': epochs,
-                # 'learning_rate': learning_rate,
-
-                # 'evaluate_data_count': evaluate_data_count,
-
-                # 'lora_r': lora_r,
-                # 'lora_alpha': lora_alpha,
-                # 'lora_dropout': lora_dropout,
-                # 'lora_target_modules': lora_target_modules,
-            }
-            if continue_from_model:
-                info['continued_from_model'] = continue_from_model
-                if continue_from_checkpoint:
-                    info['continued_from_checkpoint'] = continue_from_checkpoint
-
-            if Global.version:
-                info['tuner_version'] = Global.version
-
-            json.dump(info, info_json_file, indent=2)
-
-        if not should_training_progress_track_tqdm:
-            progress(0, desc="Train starting...")
-
-        wandb_group = template
-        wandb_tags = [f"template:{template}"]
-        if load_dataset_from == "Data Dir" and dataset_from_data_dir:
-            wandb_group += f"/{dataset_from_data_dir}"
-            wandb_tags.append(f"dataset:{dataset_from_data_dir}")
-
-        train_output = Global.finetune_train_fn(
-            base_model=base_model_name,
-            tokenizer=tokenizer_name,
-            output_dir=output_dir,
-            train_data=train_data,
-            # 128,  # batch_size (is not used, use gradient_accumulation_steps instead)
-            micro_batch_size=micro_batch_size,
-            gradient_accumulation_steps=gradient_accumulation_steps,
-            num_train_epochs=epochs,
-            learning_rate=learning_rate,
-            cutoff_len=max_seq_length,
-            val_set_size=evaluate_data_count,
-            lora_r=lora_r,
-            lora_alpha=lora_alpha,
-            lora_dropout=lora_dropout,
-            lora_target_modules=lora_target_modules,
-            lora_modules_to_save=lora_modules_to_save,
-            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_param,
-            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=Config.wandb_api_key,
-            wandb_project=Config.default_wandb_project if Config.enable_wandb else None,
-            wandb_group=wandb_group,
-            wandb_run_name=model_name,
-            wandb_tags=wandb_tags
-        )
-
-        logs_str = "\n".join([json.dumps(log)
-                             for log in log_history]) or "None"
-
-        result_message = f"Training ended:\n{str(train_output)}"
-        print(result_message)
-        # result_message += f"\n\nLogs:\n{logs_str}"
-
-        clear_cache()
-
-        return result_message
+        training_thread = threading.Thread(target=training)
+        training_thread.daemon = True
+        training_thread.start()
 
     except Exception as e:
-        raise gr.Error(
-            f"{e} (To dismiss this error, click the 'Abort' button)")
+        Global.is_training = False
+        traceback.print_exc()
+        Global.training_error_message = str(e)
+    finally:
+        Global.is_train_starting = False
+
+    return render_training_status()
+
+
+def render_training_status():
+    if not Global.is_training:
+        if Global.is_train_starting:
+            html_content = """
+            <div class="progress-block">
+              <div class="progress-level">
+                <div class="progress-level-inner">
+                  Starting...
+                </div>
+              </div>
+            </div>
+            """
+            return (gr.HTML.update(value=html_content), gr.HTML.update(visible=True))
+
+        if Global.training_error_message:
+            html_content = f"""
+            <div class="progress-block is_error">
+              <div class="progress-level">
+                <div class="error">
+                  <div class="title">
+                    ⚠ Something went wrong
+                  </div>
+                  <div class="error-message">{Global.training_error_message}</div>
+                </div>
+              </div>
+            </div>
+            """
+            return (gr.HTML.update(value=html_content), gr.HTML.update(visible=False))
+
+        if Global.train_output_str:
+            end_message = "✅ Training completed"
+            if Global.should_stop_training:
+                end_message = "🛑 Train aborted"
+            html_content = f"""
+            <div class="progress-block">
+              <div class="progress-level">
+                <div class="output">
+                  <div class="title">
+                    {end_message}
+                  </div>
+                  <div class="message">{Global.train_output_str}</div>
+                </div>
+              </div>
+            </div>
+            """
+            return (gr.HTML.update(value=html_content), gr.HTML.update(visible=False))
+
+        if Global.training_status_text:
+            html_content = f"""
+            <div class="progress-block">
+              <div class="status">{Global.training_status_text}</div>
+            </div>
+            """
+            return (gr.HTML.update(value=html_content), gr.HTML.update(visible=False))
+
+        html_content = """
+        <div class="progress-block">
+          <div class="empty-text">
+            Training status will be shown here
+          </div>
+        </div>
+        """
+        return (gr.HTML.update(value=html_content), gr.HTML.update(visible=False))
+
+    meta_info = []
+    meta_info.append(
+        f"{Global.training_current_step}/{Global.training_total_steps} steps")
+    current_time = time.time()
+    time_elapsed = current_time - Global.train_started_at
+    time_remaining = -1
+    if Global.training_eta:
+        time_remaining = Global.training_eta - current_time
+    if time_remaining >= 0:
+        meta_info.append(
+            f"{format_time(time_elapsed)}<{format_time(time_remaining)}")
+        meta_info.append(f"ETA: {format_timestamp(Global.training_eta)}")
+    else:
+        meta_info.append(format_time(time_elapsed))
+
+    html_content = f"""
+    <div class="progress-block is_training">
+      <div class="meta-text">{' | '.join(meta_info)}</div>
+      <div class="progress-level">
+        <div class="progress-level-inner">
+          {Global.training_status_text} - {Global.training_progress * 100:.2f}%
+        </div>
+        <div class="progress-bar-wrap">
+          <div class="progress-bar" style="width: {Global.training_progress * 100:.2f}%;">
+          </div>
+        </div>
+      </div>
+    </div>
+    """
+    return (gr.HTML.update(value=html_content), gr.HTML.update(visible=True))
+
+
+def format_time(seconds):
+    hours, remainder = divmod(seconds, 3600)
+    minutes, seconds = divmod(remainder, 60)
+    if hours == 0:
+        return "{:02d}:{:02d}".format(int(minutes), int(seconds))
+    else:
+        return "{:02d}:{:02d}:{:02d}".format(int(hours), int(minutes), int(seconds))
+
+
+def format_timestamp(timestamp):
+    dt_naive = datetime.datetime.utcfromtimestamp(timestamp)
+    utc = pytz.UTC
+    timezone = Config.timezone
+    dt_aware = utc.localize(dt_naive).astimezone(timezone)
+    now = datetime.datetime.now(timezone)
+    delta = dt_aware.date() - now.date()
+    if delta.days == 0:
+        time_str = ""
+    elif delta.days == 1:
+        time_str = "tomorrow at "
+    elif delta.days == -1:
+        time_str = "yesterday at "
+    else:
+        time_str = dt_aware.strftime('%A, %B %d at ')
+    time_str += dt_aware.strftime('%I:%M %p').lower()
+    return time_str

llama_lora/ui/inference_ui.py

@@ -381,7 +381,7 @@ def inference_ui():
     things_that_might_timeout = []
 
     with gr.Blocks() as inference_ui_blocks:
-        with gr.Row():
+        with gr.Row(elem_classes="disable_while_training"):
             with gr.Column(elem_id="inference_lora_model_group"):
                 model_prompt_template_message = gr.Markdown(
                     "", visible=False, elem_id="inference_lora_model_prompt_template_message")
@@ -402,7 +402,7 @@ def inference_ui():
             reload_selections_button.style(
                 full_width=False,
                 size="sm")
-        with gr.Row():
+        with gr.Row(elem_classes="disable_while_training"):
             with gr.Column():
                 with gr.Column(elem_id="inference_prompt_box"):
                     variable_0 = gr.Textbox(

llama_lora/ui/main_page.py

@@ -18,6 +18,8 @@ def main_page():
             title=title,
             css=get_css_styles(),
     ) as main_page_blocks:
+        training_indicator = gr.HTML(
+            "", visible=False, elem_id="training_indicator")
         with gr.Column(elem_id="main_page_content"):
             with gr.Row():
                 gr.Markdown(
@@ -27,7 +29,10 @@ def main_page():
                     """,
                     elem_id="page_title",
                 )
-                with gr.Column(elem_id="global_base_model_select_group"):
+                with gr.Column(
+                    elem_id="global_base_model_select_group",
+                    elem_classes="disable_while_training without_message"
+                ):
                     global_base_model_select = gr.Dropdown(
                         label="Base Model",
                         elem_id="global_base_model_select",
@@ -99,6 +104,19 @@ def main_page():
         ]
     )
 
+    main_page_blocks.load(
+        fn=lambda: gr.HTML.update(
+            visible=Global.is_training or Global.is_train_starting,
+            value=Global.is_training and "training"
+            or (
+                Global.is_train_starting and "train_starting" or ""
+            )
+        ),
+        inputs=None,
+        outputs=[training_indicator],
+        every=2
+    )
+
     main_page_blocks.load(_js=f"""
     function () {{
         {popperjs_core_code()}
@@ -239,6 +257,12 @@ def main_page_custom_css():
     }
     */
 
+   .hide_wrap > .wrap {
+       border: 0;
+       background: transparent;
+       pointer-events: none;
+   }
+
     .error-message, .error-message p {
         color: var(--error-text-color) !important;
     }
@@ -261,6 +285,36 @@ def main_page_custom_css():
         max-height: unset;
     }
 
+    #training_indicator { display: none; }
+    #training_indicator:not(.hidden) ~ * .disable_while_training {
+        position: relative !important;
+        pointer-events: none !important;
+    }
+    #training_indicator:not(.hidden) ~ * .disable_while_training * {
+        pointer-events: none !important;
+    }
+    #training_indicator:not(.hidden) ~ * .disable_while_training::after {
+        content: "Disabled while training is in progress";
+        display: flex;
+        position: absolute !important;
+        z-index: 70;
+        top: 0;
+        left: 0;
+        right: 0;
+        bottom: 0;
+        background: var(--block-background-fill);
+        opacity: 0.7;
+        justify-content: center;
+        align-items: center;
+        color: var(--body-text-color);
+        font-size: var(--text-lg);
+        font-weight: var(--weight-bold);
+        text-transform: uppercase;
+    }
+    #training_indicator:not(.hidden) ~ * .disable_while_training.without_message::after {
+        content: "";
+    }
+
     #page_title {
         flex-grow: 3;
     }

llama_lora/ui/tokenizer_ui.py

@@ -41,7 +41,7 @@ def tokenizer_ui():
     things_that_might_timeout = []
 
     with gr.Blocks() as tokenizer_ui_blocks:
-        with gr.Row():
+        with gr.Row(elem_classes="disable_while_training"):
             with gr.Column():
                 encoded_tokens = gr.Code(
                     label="Encoded Tokens (JSON)",

llama_lora/ui/trainer_callback.py

@@ -0,0 +1,104 @@
+import time
+import traceback
+from transformers import TrainerCallback
+
+from ..globals import Global
+from ..utils.eta_predictor import ETAPredictor
+
+
+def reset_training_status():
+    Global.is_train_starting = False
+    Global.is_training = False
+    Global.should_stop_training = False
+    Global.train_started_at = time.time()
+    Global.training_error_message = None
+    Global.training_error_detail = None
+    Global.training_total_epochs = 1
+    Global.training_current_epoch = 0.0
+    Global.training_total_steps = 1
+    Global.training_current_step = 0
+    Global.training_progress = 0.0
+    Global.training_log_history = []
+    Global.training_status_text = ""
+    Global.training_eta_predictor = ETAPredictor()
+    Global.training_eta = None
+    Global.train_output = None
+    Global.train_output_str = None
+
+
+def get_progress_text(current_epoch, total_epochs, last_loss):
+    progress_detail = f"Epoch {current_epoch:.2f}/{total_epochs}"
+    if last_loss is not None:
+        progress_detail += f", Loss: {last_loss:.4f}"
+    return f"Training... ({progress_detail})"
+
+
+def set_train_output(output):
+    end_by = 'aborted' if Global.should_stop_training else 'completed'
+    result_message = f"Training {end_by}"
+    Global.training_status_text = result_message
+
+    Global.train_output = output
+    Global.train_output_str = str(output)
+
+    return result_message
+
+
+def update_training_states(
+        current_step, total_steps,
+        current_epoch, total_epochs,
+        log_history):
+
+    Global.training_total_steps = total_steps
+    Global.training_current_step = current_step
+    Global.training_total_epochs = total_epochs
+    Global.training_current_epoch = current_epoch
+    Global.training_progress = current_step / total_steps
+    Global.training_log_history = log_history
+    Global.training_eta = Global.training_eta_predictor.predict_eta(current_step, total_steps)
+
+    last_history = None
+    last_loss = None
+    if len(Global.training_log_history) > 0:
+        last_history = log_history[-1]
+        last_loss = last_history.get('loss', None)
+
+    Global.training_status_text = get_progress_text(
+        total_epochs=total_epochs,
+        current_epoch=current_epoch,
+        last_loss=last_loss,
+    )
+
+
+class UiTrainerCallback(TrainerCallback):
+    def _on_progress(self, args, state, control):
+        if Global.should_stop_training:
+            control.should_training_stop = True
+
+        try:
+            total_steps = (
+                state.max_steps if state.max_steps is not None
+                else state.num_train_epochs * state.steps_per_epoch)
+            current_step = state.global_step
+
+            total_epochs = args.num_train_epochs
+            current_epoch = state.epoch
+
+            log_history = state.log_history
+
+            update_training_states(
+                total_steps=total_steps,
+                current_step=current_step,
+                total_epochs=total_epochs,
+                current_epoch=current_epoch,
+                log_history=log_history
+            )
+        except Exception as e:
+            print("Error occurred while updating UI status:", e)
+            traceback.print_exc()
+
+    def on_epoch_begin(self, args, state, control, **kwargs):
+        self._on_progress(args, state, control)
+
+    def on_step_end(self, args, state, control, **kwargs):
+        self._on_progress(args, state, control)

llama_lora/utils/eta_predictor.py

@@ -0,0 +1,54 @@
+import time
+import traceback
+from collections import deque
+from typing import Optional
+
+
+class ETAPredictor:
+    def __init__(self, lookback_minutes: int = 180):
+        self.lookback_seconds = lookback_minutes * 60  # convert minutes to seconds
+        self.data = deque()
+
+    def _cleanup_old_data(self):
+        current_time = time.time()
+        while self.data and current_time - self.data[0][1] > self.lookback_seconds:
+            self.data.popleft()
+
+    def predict_eta(
+            self, current_step: int, total_steps: int
+    ) -> Optional[int]:
+        try:
+            current_time = time.time()
+
+            # Calculate dynamic log interval based on current logged data
+            log_interval = 1
+            if len(self.data) > 100:
+                log_interval = 10
+
+            # Only log data if last log is at least log_interval seconds ago
+            if len(self.data) < 1 or current_time - self.data[-1][1] >= log_interval:
+                self.data.append((current_step, current_time))
+                self._cleanup_old_data()
+
+            # Only predict if we have enough data
+            if len(self.data) < 2 or self.data[-1][1] - self.data[0][1] < 5:
+                return None
+
+            first_step, first_time = self.data[0]
+            steps_completed = current_step - first_step
+            time_elapsed = current_time - first_time
+
+            if steps_completed == 0:
+                return None
+
+            time_per_step = time_elapsed / steps_completed
+            steps_remaining = total_steps - current_step
+
+            remaining_seconds = steps_remaining * time_per_step
+            eta_unix_timestamp = current_time + remaining_seconds
+
+            return int(eta_unix_timestamp)
+        except Exception as e:
+            print("Error predicting ETA:", e)
+            traceback.print_exc()
+            return None