update

llama_lora/globals.py

@@ -49,6 +49,7 @@ class Global:
     training_eta: Union[int, None] = None
     train_output: Union[None, Any] = None
     train_output_str: Union[None, str] = None
+    training_params_info_text: str = ""
 
     # Generation Control
     should_stop_generating: bool = False

llama_lora/lib/finetune.py

@@ -70,10 +70,12 @@ def train(
     wandb_tags: List[str] = [],
     wandb_watch: str = "false",  # options: false | gradients | all
     wandb_log_model: str = "true",  # options: false | true
+    additional_wandb_config: Union[dict, None] = None,
     status_message_callback: Any = None,
+    params_info_callback: Any = None,
 ):
     if status_message_callback:
-        cb_result = status_message_callback("Preparing training...")
+        cb_result = status_message_callback("Preparing...")
         if cb_result:
             return
 
@@ -163,6 +165,8 @@ def train(
             config={'finetune_args': finetune_args},
             # id=None  # used for resuming
         )
+        if additional_wandb_config:
+            wandb.config.update(additional_wandb_config)
     else:
         os.environ['WANDB_MODE'] = "disabled"
 
@@ -294,6 +298,10 @@ def train(
     if use_wandb and wandb:
         wandb.config.update({"model": {"all_params": all_params, "trainable_params": trainable_params,
                             "trainable%": 100 * trainable_params / all_params}})
+    if params_info_callback:
+        cb_result = params_info_callback(all_params=all_params, trainable_params=trainable_params)
+        if cb_result:
+            return
 
     if status_message_callback:
         cb_result = status_message_callback("Preparing train data...")

llama_lora/ui/finetune/style.css

@@ -271,6 +271,7 @@
 #finetune_training_status .progress-block {
     min-height: 100px;
     display: flex;
+    flex-direction: column;
     justify-content: center;
     align-items: center;
     background: var(--panel-background-fill);
@@ -300,11 +301,14 @@
     white-space: pre-wrap;
 }
 #finetune_training_status .progress-block .progress-level {
+    flex-grow: 1;
     display: flex;
     flex-direction: column;
+    justify-content: center;
     align-items: center;
     z-index: var(--layer-2);
     width: var(--size-full);
+    padding: 8px 0;
 }
 #finetune_training_status .progress-block .progress-level-inner {
     margin: var(--size-2) auto;
@@ -326,6 +330,17 @@
     transition: all 150ms ease 0s;
 }
 
+#finetune_training_status .progress-block .params-info {
+    font-size: var(--text-sm);
+    font-weight: var(--weight-light);
+    margin-top: 8px;
+    margin-bottom: -4px !important;
+    opacity: 0.4;
+}
+#finetune_training_status .progress-block .progress-level + .params-info {
+    margin-top: -8px;
+}
+
 #finetune_training_status .progress-block .output {
     display: flex;
     flex-direction: column;

llama_lora/ui/finetune/training.py

@@ -29,6 +29,10 @@ def status_message_callback(message):
     Global.training_status_text = message
 
 
+def params_info_callback(all_params, trainable_params):
+    Global.training_params_info_text = f"Params: {trainable_params}/{all_params} ({100 * trainable_params / all_params:.4f}% trainable)"
+
+
 def do_train(
     # Dataset
     template,
@@ -262,6 +266,8 @@ def do_train(
                     train_data=train_data,
                     callbacks=training_callbacks,
                     status_message_callback=status_message_callback,
+                    params_info_callback=params_info_callback,
+                    additional_wandb_config=info,
                     **finetune_args,
                 )
 
@@ -325,6 +331,14 @@ def render_training_status():
             end_message = "✅ Training completed"
             if Global.should_stop_training:
                 end_message = "🛑 Train aborted"
+
+            params_info_html = ""
+            if Global.training_params_info_text:
+                params_info_html = f"""
+                <div class="params-info">
+                  {Global.training_params_info_text}
+                </div>
+                """
             html_content = f"""
             <div class="progress-block">
               <div class="progress-level">
@@ -335,6 +349,7 @@ def render_training_status():
                   <div class="message">{Global.train_output_str}</div>
                 </div>
               </div>
+              {params_info_html}
             </div>
             """
             return (gr.HTML.update(value=html_content), gr.HTML.update(visible=False))
@@ -371,6 +386,13 @@ def render_training_status():
     else:
         meta_info.append(format_time(time_elapsed))
 
+    params_info_html = ""
+    if Global.training_params_info_text:
+        params_info_html = f"""
+        <div class="params-info">
+          {Global.training_params_info_text}
+        </div>
+        """
     html_content = f"""
     <div class="progress-block is_training">
       <div class="meta-text">{' | '.join(meta_info)}</div>
@@ -383,6 +405,7 @@ def render_training_status():
           </div>
         </div>
       </div>
+      {params_info_html}
     </div>
     """
     return (gr.HTML.update(value=html_content), gr.HTML.update(visible=True))

llama_lora/ui/trainer_callback.py

@@ -24,6 +24,7 @@ def reset_training_status():
     Global.training_eta = None
     Global.train_output = None
     Global.train_output_str = None
+    Global.training_params_info_text = ""
 
 
 def get_progress_text(current_epoch, total_epochs, last_loss):