Multiple finetune models (#179)

* add finetune info to models tab

* show current completion model in finetune tab

* modal warning if checkpoint do not match with selected model

* Revert "Revert "hidden => False for starcoder models (#170)""

This reverts commit 52203da.

* fix model name for lora setup

* fix table styles

Author: mitya52

known_models_db/refact_known_models/huggingface.py

@@ -31,35 +31,35 @@
     },
     "starcoder/1b/base": {
         "backend": "transformers",
-        "model_path": "bigcode/starcoderbase-1b",
+        "model_path": "smallcloudai/starcoderbase-1b",
         "diff_scratchpad_class": "refact_scratchpads:ScratchpadPSM",
         "chat_scratchpad_class": None,
         "model_class_kwargs": {},
         "required_memory_mb": 6000,
         "T": 4096,
-        "hidden": True,
+        "hidden": False,
         "filter_caps": ["completion", "finetune"],
     },
     "starcoder/3b/base": {
         "backend": "transformers",
-        "model_path": "bigcode/starcoderbase-3b",
+        "model_path": "smallcloudai/starcoderbase-3b",
         "diff_scratchpad_class": "refact_scratchpads:ScratchpadPSM",
         "chat_scratchpad_class": None,
         "model_class_kwargs": {},
         "required_memory_mb": 9000,
         "T": 4096,
-        "hidden": True,
+        "hidden": False,
         "filter_caps": ["completion", "finetune"],
     },
     "starcoder/7b/base": {
         "backend": "transformers",
-        "model_path": "bigcode/starcoderbase-7b",
+        "model_path": "smallcloudai/starcoderbase-7b",
         "diff_scratchpad_class": "refact_scratchpads:ScratchpadPSM",
         "chat_scratchpad_class": None,
         "model_class_kwargs": {},
         "required_memory_mb": 18000,
         "T": 2048,
-        "hidden": True,
+        "hidden": False,
         "filter_caps": ["completion", "finetune"],
     },
     "wizardcoder/15b": {

self_hosting_machinery/webgui/selfhost_model_assigner.py

@@ -8,6 +8,8 @@
 from self_hosting_machinery.webgui.selfhost_webutils import log
 from known_models_db.refact_known_models import models_mini_db
 from known_models_db.refact_toolbox_db import modelcap_records
+from self_hosting_machinery.scripts.best_lora import find_best_lora
+from refact_data_pipeline.finetune.finetune_utils import get_active_loras
 
 from typing import List, Dict, Set, Any
 
@@ -205,12 +207,28 @@ def _capabilities(func_type: str) -> Set:
 
         chat_caps = _capabilities("chat")
         toolbox_caps = _capabilities("toolbox")
+        active_loras = get_active_loras(self.models_db)
         for k, rec in self.models_db.items():
             if rec.get("hidden", False):
                 continue
+            finetune_info = None
+            if k in active_loras:
+                lora_mode = active_loras[k]["lora_mode"]
+                latest_best_lora_info = find_best_lora(k)
+                if lora_mode == "latest-best" and latest_best_lora_info["latest_run_id"]:
+                    finetune_info = {
+                        "run": latest_best_lora_info["latest_run_id"],
+                        "checkpoint": latest_best_lora_info["best_checkpoint_id"],
+                    }
+                elif lora_mode == "specific" and active_loras[k].get("specific_lora_run_id", ""):
+                    finetune_info = {
+                        "run": active_loras[k]["specific_lora_run_id"],
+                        "checkpoint": active_loras[k]["specific_checkpoint"],
+                    }
             info.append({
                 "name": k,
                 "backend": rec["backend"],
+                "finetune_info": finetune_info,
                 "has_completion": bool("completion" in rec["filter_caps"]),
                 "has_finetune": bool("finetune" in rec["filter_caps"]),
                 "has_toolbox": bool(toolbox_caps.intersection(rec["filter_caps"])),

self_hosting_machinery/webgui/static/tab-finetune.html

@@ -4,6 +4,7 @@
   <div class="col-5">
     <div class="pane use-model-pane">
       <h3>Use Finetuned Model</h3>
+      <div class="lora-model"><div id="lora-switch-model">Model:</div></div>
       <div class="lora-group">
         <div class="btn-group" role="group" aria-label="basic radio toggle button group">
           <input type="radio" class="lora-switch btn-check" name="finetune_lora" value="off" id="loraradio1"
@@ -387,4 +388,23 @@ <h5>Limit training time</h5>
     </div>
   </div>
 
+  <div class="modal fade" id="finetune-tab-model-warning-modal" tabindex="-1" aria-labelledby="finetune-tab-invalid-model-modal" aria-hidden="true">
+    <div class="modal-dialog modal-lg modal-dialog-centered">
+      <div class="modal-content">
+        <div class="modal-header">
+          <h5 class="modal-title" id="urlModalLabel">Warning</h5>
+          <button type="button" class="btn-close" data-bs-dismiss="modal" aria-label="Close"></button>
+        </div>
+        <div class="modal-body">
+          <div class="row">
+            <div class="mb-3" id="model-warning-message"></div>
+          </div>
+        </div>
+        <div class="modal-footer">
+          <button type="button" class="btn btn-secondary" data-bs-dismiss="modal">Ok</button>
+        </div>
+      </div>
+    </div>
+  </div>
+
 </div>

self_hosting_machinery/webgui/static/tab-finetune.js

@@ -54,7 +54,6 @@ function tab_finetune_config_and_runs() {
             return response.json();
         })
         .then(function (data) {
-            console.log('tab-finetune-config-and-runs',data);
             finetune_configs_and_runs = data;
             render_runs();
             render_model_select();
@@ -274,7 +273,12 @@ const find_checkpoints_by_run = (run_id) => {
 };
 
 function render_lora_switch() {
-    let mode = finetune_configs_and_runs.active[finetune_configs_and_runs.config.model_name] ? finetune_configs_and_runs.active[finetune_configs_and_runs.config.model_name].lora_mode : "latest-best";
+    const model_name = finetune_configs_and_runs.completion_model;
+    let lora_switch_model = document.querySelector('#lora-switch-model');
+    lora_switch_model.innerHTML = `
+        <b>Model:</b> ${model_name}
+    `;
+    let mode = finetune_configs_and_runs.active[model_name] ? finetune_configs_and_runs.active[model_name].lora_mode : "latest-best";
     loras_switch_no_reaction = true; // avoid infinite loop when setting .checked
     if (mode === 'off') {
         loras_switch_off.checked = true;
@@ -291,8 +295,8 @@ function render_lora_switch() {
         lora_switch_checkpoint.style.display = 'block';
         lora_switch_run_id.style.opacity = 1;
         lora_switch_checkpoint.style.opacity = 1;
-        lora_switch_run_id.innerHTML = `<b>Run:</b> ${finetune_configs_and_runs.active[finetune_configs_and_runs.config.model_name].specific_lora_run_id}`;
-        lora_switch_checkpoint.innerHTML = `<b>Checkpoint:</b> ${finetune_configs_and_runs.active[finetune_configs_and_runs.config.model_name].specific_checkpoint}`;
+        lora_switch_run_id.innerHTML = `<b>Run:</b> ${finetune_configs_and_runs.active[model_name].specific_lora_run_id}`;
+        lora_switch_checkpoint.innerHTML = `<b>Checkpoint:</b> ${finetune_configs_and_runs.active[model_name].specific_checkpoint}`;
     } else if (mode == 'latest-best') {
         lora_switch_run_id.style.display = 'block';
         lora_switch_checkpoint.style.display = 'block';
@@ -303,8 +307,8 @@ function render_lora_switch() {
     } else {
         lora_switch_run_id.style.display = 'none';
         lora_switch_checkpoint.style.display = 'none';
-        lora_switch_run_id.innerHTML = `<b>Run:</b> ${finetune_configs_and_runs.active[finetune_configs_and_runs.config.model_name].specific_lora_run_id}`;
-        lora_switch_checkpoint.innerHTML = `<b>Checkpoint:</b> ${finetune_configs_and_runs.active[finetune_configs_and_runs.config.model_name].specific_checkpoint}`;
+        lora_switch_run_id.innerHTML = `<b>Run:</b> ${finetune_configs_and_runs.active[model_name].specific_lora_run_id}`;
+        lora_switch_checkpoint.innerHTML = `<b>Checkpoint:</b> ${finetune_configs_and_runs.active[model_name].specific_checkpoint}`;
     }
 }
 
@@ -346,7 +350,20 @@ function render_checkpoints(data = []) {
                     }
                     row.classList.add('table-success');
                 }
-                finetune_switch_activate("specific", selected_lora, cell.dataset.checkpoint);
+                const finetune_run = finetune_configs_and_runs.finetune_runs.find((run) => run.run_id === selected_lora);
+                if (finetune_run && finetune_run.model_name !== finetune_configs_and_runs.completion_model) {
+                    let modal = document.getElementById('finetune-tab-model-warning-modal');
+                    let modal_instance = bootstrap.Modal.getOrCreateInstance(modal);
+                    document.querySelector('#finetune-tab-model-warning-modal #model-warning-message').innerHTML = `
+                    <label>
+                        Checkpoint you're about to activate trained for <b>${finetune_run.model_name}</b> model.
+                        Use another checkpoint for <b>${finetune_configs_and_runs.completion_model}</b> model instead.
+                    </label>
+                    `;
+                    modal_instance.show();
+                } else {
+                    finetune_switch_activate("specific", selected_lora, cell.dataset.checkpoint);
+                }
             });
         });
     }
@@ -361,7 +378,7 @@ function animate_use_model() {
 
 function finetune_switch_activate(lora_mode, run_id, checkpoint) {
     animate_use_model();
-    const model_name = document.querySelector('#finetune-model').value
+    const model_name = finetune_configs_and_runs.completion_model;
     let send_this = {
         "model": model_name,
         "lora_mode": lora_mode,
@@ -876,6 +893,8 @@ function start_log_stream(run_id) {
     };
     fetchData();
 }
+
+
 export async function init() {
     let req = await fetch('/tab-finetune.html');
     document.querySelector('#finetune').innerHTML = await req.text();

self_hosting_machinery/webgui/static/tab-model-hosting.html

@@ -9,6 +9,7 @@ <h3>Hosted Models</h3>
       <tr>
         <th>Model</th>
         <th>Completion</th>
+        <th>Finetune</th>
         <th>Sharding</th>
         <th>Share GPU</th>
         <th></th>

self_hosting_machinery/webgui/static/tab-model-hosting.js

@@ -146,6 +146,7 @@ function render_models_assigned(models) {
         row.setAttribute('data-model',index);
         const model_name = document.createElement("td");
         const completion = document.createElement("td");
+        const finetune_info = document.createElement("td");
         const select_gpus = document.createElement("td");
         const gpus_share = document.createElement("td");
         const del = document.createElement("td");
@@ -168,6 +169,21 @@ function render_models_assigned(models) {
             completion.appendChild(completion_input);
         }
 
+        if (models_info[index].hasOwnProperty('finetune_info') && models_info[index].finetune_info) {
+            finetune_info.innerHTML = `
+            <table cellpadding="5">
+                <tr>
+                    <td>Run: </td>
+                    <td>${models_info[index].finetune_info.run}</td>
+                </tr>
+                <tr>
+                    <td>Checkpoint: </td>
+                    <td>${models_info[index].finetune_info.checkpoint}</td>
+                </tr>
+            </table>
+            `;
+        }
+
          if (models_info[index].hasOwnProperty('has_sharding') && models_info[index].has_sharding) {
             const select_gpus_div = document.createElement("div");
             select_gpus_div.setAttribute("class", "btn-group btn-group-sm");
@@ -233,6 +249,7 @@ function render_models_assigned(models) {
 
         row.appendChild(model_name);
         row.appendChild(completion);
+        row.appendChild(finetune_info);
         row.appendChild(select_gpus);
         row.appendChild(gpus_share);
         row.appendChild(del);

self_hosting_machinery/webgui/tab_finetune.py

@@ -10,6 +10,7 @@
 from fastapi.responses import Response, StreamingResponse, JSONResponse
 
 from self_hosting_machinery.scripts import best_lora
+from self_hosting_machinery.webgui.selfhost_model_assigner import ModelAssigner
 from refact_data_pipeline.finetune.finetune_utils import get_active_loras
 from refact_data_pipeline.finetune.finetune_utils import get_finetune_config
 from refact_data_pipeline.finetune.finetune_utils import get_finetune_filter_stat
@@ -93,7 +94,7 @@ class TabFinetuneTrainingSetup(BaseModel):
 
 class TabFinetuneRouter(APIRouter):
 
-    def __init__(self, models_db: Dict[str, Any], *args, **kwargs):
+    def __init__(self, model_assigner: ModelAssigner, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.add_api_route("/tab-finetune-get", self._tab_finetune_get, methods=["GET"])
         self.add_api_route("/tab-finetune-config-and-runs", self._tab_finetune_config_and_runs, methods=["GET"])
@@ -109,7 +110,7 @@ def __init__(self, models_db: Dict[str, Any], *args, **kwargs):
         self.add_api_route("/tab-finetune-smart-filter-get", self._tab_finetune_smart_filter_get, methods=["GET"])
         self.add_api_route("/tab-finetune-training-setup", self._tab_finetune_training_setup, methods=["POST"])
         self.add_api_route("/tab-finetune-training-get", self._tab_finetune_training_get, methods=["GET"])
-        self._models_db = models_db
+        self._model_assigner = model_assigner
 
     async def _tab_finetune_get(self):
         prog, status = get_prog_and_status_for_ui()
@@ -140,9 +141,11 @@ async def _tab_finetune_get_sources_status(self):
             return f"Error: {str(e)}"
 
     async def _tab_finetune_config_and_runs(self):
+        completion_model = self._model_assigner.model_assignment.get("completion", "")
         runs = get_finetune_runs()
-        config = get_finetune_config(self._models_db)
+        config = get_finetune_config(self._model_assigner.models_db)
         result = {
+            "completion_model": completion_model,
             "finetune_runs": runs,
             "config": {
                 "limit_training_time_minutes": "60",
@@ -151,8 +154,8 @@ async def _tab_finetune_config_and_runs(self):
                 "auto_delete_n_runs": "5",
                 **config,  # TODO: why we mix finetune config for training and schedule?
             },
-            "active": get_active_loras(self._models_db),
-            "finetune_latest_best": best_lora.find_best_lora(config["model_name"]),
+            "active": get_active_loras(self._model_assigner.models_db),
+            "finetune_latest_best": best_lora.find_best_lora(completion_model),
         }
         return Response(json.dumps(result, indent=4) + "\n")
 
@@ -188,7 +191,7 @@ async def _tab_finetune_training_setup(self, post: TabFinetuneTrainingSetup):
     async def _tab_finetune_training_get(self):
         result = {
             "defaults": finetune_train_defaults,
-            "user_config": get_finetune_config(self._models_db),
+            "user_config": get_finetune_config(self._model_assigner.models_db),
         }
         return Response(json.dumps(result, indent=4) + "\n")
 
@@ -255,7 +258,7 @@ async def _tab_finetune_remove(self, run_id: str):
         return JSONResponse("OK")
 
     async def _tab_finetune_activate(self, activate: TabFinetuneActivate):
-        active_loras = get_active_loras(self._models_db)
+        active_loras = get_active_loras(self._model_assigner.models_db)
         active_loras[activate.model] = activate.dict()
         with open(env.CONFIG_ACTIVE_LORA, "w") as f:
             json.dump(active_loras, f, indent=4)

self_hosting_machinery/webgui/webgui.py

@@ -75,7 +75,7 @@ def _routers_list(
             TabServerLogRouter(),
             TabUploadRouter(),
             TabFinetuneRouter(
-                models_db=model_assigner.models_db),
+                model_assigner=model_assigner),
             TabHostRouter(model_assigner),
             TabSettingsRouter(model_assigner),
             StaticRouter(),