change model loading mechanism

LLaMA_LoRA.ipynb

@@ -281,7 +281,7 @@
         "\n",
         "# Set Configs\n",
         "from llama_lora.llama_lora.globals import Global\n",
-        "Global.base_model = base_model\n",
+        "Global.default_base_model_name = base_model\n",
         "data_dir_realpath = !realpath ./data\n",
         "Global.data_dir = data_dir_realpath[0]\n",
         "Global.load_8bit = True\n",
@@ -289,12 +289,7 @@
         "# Prepare Data Dir\n",
         "import os\n",
         "from llama_lora.llama_lora.utils.data import init_data_dir\n",
-        "init_data_dir()\n",
-        "\n",
-        "# Load the Base Model\n",
-        "from llama_lora.llama_lora.models import load_base_model\n",
-        "load_base_model()\n",
-        "print(f\"Base model loaded: '{Global.base_model}'.\")"
+        "init_data_dir()"
       ],
       "metadata": {
         "id": "Yf6g248ylteP"

app.py

@@ -7,7 +7,6 @@ import gradio as gr
 from llama_lora.globals import Global
 from llama_lora.ui.main_page import main_page, get_page_title, main_page_custom_css
 from llama_lora.utils.data import init_data_dir
-from llama_lora.models import load_base_model
 
 
 def main(
@@ -31,7 +30,7 @@ def main(
         data_dir
     ), "Please specify a --data_dir, e.g. --data_dir='./data'"
 
-    Global.base_model = base_model
+    Global.default_base_model_name = base_model
     Global.data_dir = os.path.abspath(data_dir)
     Global.load_8bit = load_8bit
 
@@ -41,9 +40,6 @@ def main(
     os.makedirs(data_dir, exist_ok=True)
     init_data_dir()
 
-    if not skip_loading_base_model:
-        load_base_model()
-
     with gr.Blocks(title=get_page_title(), css=main_page_custom_css()) as demo:
         main_page()
 

llama_lora/globals.py

@@ -13,12 +13,10 @@ from .lib.finetune import train
 class Global:
     version = None
 
-    base_model: str = ""
     data_dir: str = ""
     load_8bit: bool = False
 
-    loaded_tokenizer: Any = None
-    loaded_base_model: Any = None
+    default_base_model_name: str = ""
 
     # Functions
     train_fn: Any = train
@@ -31,8 +29,8 @@ class Global:
     generation_force_stopped_at = None
 
     # Model related
-    model_has_been_used = False
-    cached_lora_models = LRUCache(10)
+    loaded_models = LRUCache(1)
+    loaded_tokenizers = LRUCache(1)
 
     # GPU Info
     gpu_cc = None  # GPU compute capability

llama_lora/models.py

@@ -3,9 +3,8 @@ import sys
 import gc
 
 import torch
-import transformers
+from transformers import LlamaForCausalLM, LlamaTokenizer
 from peft import PeftModel
-from transformers import GenerationConfig, LlamaForCausalLM, LlamaTokenizer
 
 from .globals import Global
 
@@ -23,96 +22,120 @@ def get_device():
         pass
 
 
-device = get_device()
-
-
-def get_base_model():
-    load_base_model()
-    return Global.loaded_base_model
-
-
-def get_model_with_lora(lora_weights_name_or_path: str = "tloen/alpaca-lora-7b"):
-    Global.model_has_been_used = True
+def get_new_base_model(base_model_name):
+    if Global.ui_dev_mode:
+        return
 
-    if Global.cached_lora_models:
-        model_from_cache = Global.cached_lora_models.get(lora_weights_name_or_path)
-        if model_from_cache:
-            return model_from_cache
+    device = get_device()
 
     if device == "cuda":
-        model = PeftModel.from_pretrained(
-            get_base_model(),
-            lora_weights_name_or_path,
+        model = LlamaForCausalLM.from_pretrained(
+            base_model_name,
+            load_in_8bit=Global.load_8bit,
             torch_dtype=torch.float16,
+            # device_map="auto",
             device_map={'': 0},  # ? https://github.com/tloen/alpaca-lora/issues/21
         )
     elif device == "mps":
-        model = PeftModel.from_pretrained(
-            get_base_model(),
-            lora_weights_name_or_path,
+        model = LlamaForCausalLM.from_pretrained(
+            base_model_name,
             device_map={"": device},
             torch_dtype=torch.float16,
         )
     else:
-        model = PeftModel.from_pretrained(
-            get_base_model(),
-            lora_weights_name_or_path,
-            device_map={"": device},
+        model = LlamaForCausalLM.from_pretrained(
+            base_model_name, device_map={"": device}, low_cpu_mem_usage=True
         )
 
-    model.config.pad_token_id = get_tokenizer().pad_token_id = 0
+    model.config.pad_token_id = get_tokenizer(base_model_name).pad_token_id = 0
     model.config.bos_token_id = 1
     model.config.eos_token_id = 2
 
-    if not Global.load_8bit:
-        model.half()  # seems to fix bugs for some users.
+    return model
 
-    model.eval()
-    if torch.__version__ >= "2" and sys.platform != "win32":
-        model = torch.compile(model)
 
-    if Global.cached_lora_models:
-        Global.cached_lora_models.set(lora_weights_name_or_path, model)
+def get_tokenizer(base_model_name):
+    if Global.ui_dev_mode:
+        return
 
-    return model
+    loaded_tokenizer = Global.loaded_tokenizers.get(base_model_name)
+    if loaded_tokenizer:
+        return loaded_tokenizer
 
+    tokenizer = LlamaTokenizer.from_pretrained(base_model_name)
+    Global.loaded_tokenizers.set(base_model_name, tokenizer)
 
-def get_tokenizer():
-    load_base_model()
-    return Global.loaded_tokenizer
+    return tokenizer
 
 
-def load_base_model():
+def get_model(
+    base_model_name,
+    peft_model_name = None):
     if Global.ui_dev_mode:
         return
 
-    if Global.loaded_tokenizer is None:
-        Global.loaded_tokenizer = LlamaTokenizer.from_pretrained(
-            Global.base_model
-        )
-    if Global.loaded_base_model is None:
+    if peft_model_name == "None":
+        peft_model_name = None
+
+    model_key = base_model_name
+    if peft_model_name:
+        model_key = f"{base_model_name}//{peft_model_name}"
+
+    loaded_model = Global.loaded_models.get(model_key)
+    if loaded_model:
+        return loaded_model
+
+    peft_model_name_or_path = peft_model_name
+
+    lora_models_directory_path = os.path.join(Global.data_dir, "lora_models")
+    possible_lora_model_path = os.path.join(lora_models_directory_path, peft_model_name)
+    if os.path.isdir(possible_lora_model_path):
+        peft_model_name_or_path = possible_lora_model_path
+
+    Global.loaded_models.prepare_to_set()
+    clear_cache()
+
+    model = get_new_base_model(base_model_name)
+
+    if peft_model_name:
+        device = get_device()
+
         if device == "cuda":
-            Global.loaded_base_model = LlamaForCausalLM.from_pretrained(
-                Global.base_model,
-                load_in_8bit=Global.load_8bit,
+            model = PeftModel.from_pretrained(
+                model,
+                peft_model_name_or_path,
                 torch_dtype=torch.float16,
-                # device_map="auto",
                 device_map={'': 0},  # ? https://github.com/tloen/alpaca-lora/issues/21
             )
         elif device == "mps":
-            Global.loaded_base_model = LlamaForCausalLM.from_pretrained(
-                Global.base_model,
+            model = PeftModel.from_pretrained(
+                model,
+                peft_model_name_or_path,
                 device_map={"": device},
                 torch_dtype=torch.float16,
             )
         else:
-            Global.loaded_base_model = LlamaForCausalLM.from_pretrained(
-                Global.base_model, device_map={"": device}, low_cpu_mem_usage=True
+            model = PeftModel.from_pretrained(
+                model,
+                peft_model_name_or_path,
+                device_map={"": device},
             )
 
-        Global.loaded_base_model.config.pad_token_id = get_tokenizer().pad_token_id = 0
-        Global.loaded_base_model.config.bos_token_id = 1
-        Global.loaded_base_model.config.eos_token_id = 2
+    model.config.pad_token_id = get_tokenizer(base_model_name).pad_token_id = 0
+    model.config.bos_token_id = 1
+    model.config.eos_token_id = 2
+
+    if not Global.load_8bit:
+        model.half()  # seems to fix bugs for some users.
+
+    model.eval()
+    if torch.__version__ >= "2" and sys.platform != "win32":
+        model = torch.compile(model)
+
+    Global.loaded_models.set(model_key, model)
+    clear_cache()
+
+    return model
 
 
 def clear_cache():
@@ -124,19 +147,137 @@ def clear_cache():
 
 
 def unload_models():
-    del Global.loaded_base_model
-    Global.loaded_base_model = None
+    Global.loaded_models.clear()
+    Global.loaded_tokenizers.clear()
+    clear_cache()
 
-    del Global.loaded_tokenizer
-    Global.loaded_tokenizer = None
 
-    Global.cached_lora_models.clear()
 
-    clear_cache()
 
-    Global.model_has_been_used = False
+
+########
+
+# def get_base_model():
+#     load_base_model()
+#     return Global.loaded_base_model
+
+
+# def get_model_with_lora(lora_weights_name_or_path: str = "tloen/alpaca-lora-7b"):
+#     # Global.model_has_been_used = True
+#     #
+#     #
+#     if Global.loaded_tokenizer is None:
+#         Global.loaded_tokenizer = LlamaTokenizer.from_pretrained(
+#             Global.base_model
+#         )
+
+#     if Global.cached_lora_models:
+#         model_from_cache = Global.cached_lora_models.get(lora_weights_name_or_path)
+#         if model_from_cache:
+#             return model_from_cache
+
+#     Global.cached_lora_models.prepare_to_set()
+
+#     if device == "cuda":
+#         model = PeftModel.from_pretrained(
+#             get_new_base_model(),
+#             lora_weights_name_or_path,
+#             torch_dtype=torch.float16,
+#             device_map={'': 0},  # ? https://github.com/tloen/alpaca-lora/issues/21
+#         )
+#     elif device == "mps":
+#         model = PeftModel.from_pretrained(
+#             get_new_base_model(),
+#             lora_weights_name_or_path,
+#             device_map={"": device},
+#             torch_dtype=torch.float16,
+#         )
+#     else:
+#         model = PeftModel.from_pretrained(
+#             get_new_base_model(),
+#             lora_weights_name_or_path,
+#             device_map={"": device},
+#         )
+
+#     model.config.pad_token_id = get_tokenizer().pad_token_id = 0
+#     model.config.bos_token_id = 1
+#     model.config.eos_token_id = 2
+
+#     if not Global.load_8bit:
+#         model.half()  # seems to fix bugs for some users.
+
+#     model.eval()
+#     if torch.__version__ >= "2" and sys.platform != "win32":
+#         model = torch.compile(model)
+
+#     if Global.cached_lora_models:
+#         Global.cached_lora_models.set(lora_weights_name_or_path, model)
+
+#     clear_cache()
+
+#     return model
+
+
+
+
+
+# def load_base_model():
+#     return;
+
+#     if Global.ui_dev_mode:
+#         return
+
+#     if Global.loaded_tokenizer is None:
+#         Global.loaded_tokenizer = LlamaTokenizer.from_pretrained(
+#             Global.base_model
+#         )
+#     if Global.loaded_base_model is None:
+#         if device == "cuda":
+#             Global.loaded_base_model = LlamaForCausalLM.from_pretrained(
+#                 Global.base_model,
+#                 load_in_8bit=Global.load_8bit,
+#                 torch_dtype=torch.float16,
+#                 # device_map="auto",
+#                 device_map={'': 0},  # ? https://github.com/tloen/alpaca-lora/issues/21
+#             )
+#         elif device == "mps":
+#             Global.loaded_base_model = LlamaForCausalLM.from_pretrained(
+#                 Global.base_model,
+#                 device_map={"": device},
+#                 torch_dtype=torch.float16,
+#             )
+#         else:
+#             Global.loaded_base_model = LlamaForCausalLM.from_pretrained(
+#                 Global.base_model, device_map={"": device}, low_cpu_mem_usage=True
+#             )
+
+#         Global.loaded_base_model.config.pad_token_id = get_tokenizer().pad_token_id = 0
+#         Global.loaded_base_model.config.bos_token_id = 1
+#         Global.loaded_base_model.config.eos_token_id = 2
+
+
+# def clear_cache():
+#     gc.collect()
+
+#     # if not shared.args.cpu: # will not be running on CPUs anyway
+#     with torch.no_grad():
+#         torch.cuda.empty_cache()
+
+
+# def unload_models():
+#     del Global.loaded_base_model
+#     Global.loaded_base_model = None
+
+#     del Global.loaded_tokenizer
+#     Global.loaded_tokenizer = None
+
+#     Global.cached_lora_models.clear()
+
+#     clear_cache()
+
+#     Global.model_has_been_used = False
 
 
-def unload_models_if_already_used():
-    if Global.model_has_been_used:
-        unload_models()
+# def unload_models_if_already_used():
+#     if Global.model_has_been_used:
+#         unload_models()

llama_lora/ui/finetune_ui.py

@@ -10,8 +10,8 @@ from transformers import TrainerCallback
 
 from ..globals import Global
 from ..models import (
-    get_base_model, get_tokenizer,
-    clear_cache, unload_models_if_already_used)
+    get_new_base_model, get_tokenizer,
+    clear_cache, unload_models)
 from ..utils.data import (
     get_available_template_names,
     get_available_dataset_names,
@@ -269,14 +269,16 @@ def do_train(
     progress=gr.Progress(track_tqdm=should_training_progress_track_tqdm),
 ):
     try:
+        base_model_name = Global.default_base_model_name
+        output_dir = os.path.join(Global.data_dir, "lora_models", model_name)
+        if os.path.exists(output_dir):
+            if (not os.path.isdir(output_dir)) or os.path.exists(os.path.join(output_dir, 'adapter_config.json')):
+                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...")
 
-        clear_cache()
-        # If model has been used in inference, we need to unload it first.
-        # Otherwise, we'll get a 'Function MmBackward0 returned an invalid
-        # gradient at index 1 - expected device meta but got cuda:0' error.
-        unload_models_if_already_used()
+        unload_models()  # Need RAM for training
 
         prompter = Prompter(template)
         variable_names = prompter.get_variable_names()
@@ -415,17 +417,12 @@ Train data (first 10):
 
         Global.should_stop_training = False
 
-        # Do this again right before training to make sure the model is not used in inference.
-        unload_models_if_already_used()
-        clear_cache()
-
-        base_model = get_base_model()
-        tokenizer = get_tokenizer()
+        base_model = get_new_base_model(base_model_name)
+        tokenizer = get_tokenizer(base_model_name)
 
         # 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.
 
-        output_dir = os.path.join(Global.data_dir, "lora_models", model_name)
         if not os.path.exists(output_dir):
             os.makedirs(output_dir)
 
@@ -435,10 +432,11 @@ Train data (first 10):
                 dataset_name = dataset_from_data_dir
 
             info = {
-                'base_model': Global.base_model,
+                'base_model': base_model_name,
                 'prompt_template': template,
                 'dataset_name': dataset_name,
                 'dataset_rows': len(train_data),
+                'timestamp': time.time()
             }
             json.dump(info, info_json_file, indent=2)
 
@@ -472,7 +470,11 @@ Train data (first 10):
 
         result_message = f"Training ended:\n{str(train_output)}\n\nLogs:\n{logs_str}"
         print(result_message)
+
+        del base_model
+        del tokenizer
         clear_cache()
+
         return result_message
 
     except Exception as e:
@@ -837,6 +839,12 @@ def finetune_ui():
                 document.getElementById('finetune_confirm_stop_btn').style.display =
                   'none';
               }, 5000);
+              document.getElementById('finetune_confirm_stop_btn').style['pointer-events'] =
+                'none';
+              setTimeout(function () {
+                document.getElementById('finetune_confirm_stop_btn').style['pointer-events'] =
+                  'inherit';
+              }, 300);
               document.getElementById('finetune_stop_btn').style.display = 'none';
               document.getElementById('finetune_confirm_stop_btn').style.display =
                 'block';

llama_lora/ui/inference_ui.py

@@ -7,11 +7,10 @@ import transformers
 from transformers import GenerationConfig
 
 from ..globals import Global
-from ..models import get_base_model, get_model_with_lora, get_tokenizer, get_device
+from ..models import get_model, get_tokenizer, get_device
 from ..utils.data import (
     get_available_template_names,
     get_available_lora_model_names,
-    get_path_of_available_lora_model,
     get_info_of_available_lora_model)
 from ..utils.prompter import Prompter
 from ..utils.callbacks import Iteratorize, Stream
@@ -22,6 +21,18 @@ default_show_raw = True
 inference_output_lines = 12
 
 
+def prepare_inference(lora_model_name, progress=gr.Progress(track_tqdm=True)):
+    base_model_name = Global.default_base_model_name
+
+    try:
+        get_tokenizer(base_model_name)
+        get_model(base_model_name, lora_model_name)
+        return ("", "")
+
+    except Exception as e:
+        raise gr.Error(e)
+
+
 def do_inference(
     lora_model_name,
     prompt_template,
@@ -37,6 +48,8 @@ def do_inference(
     show_raw=False,
     progress=gr.Progress(track_tqdm=True),
 ):
+    base_model_name = Global.default_base_model_name
+
     try:
         if Global.generation_force_stopped_at is not None:
             required_elapsed_time_after_forced_stop = 1
@@ -52,16 +65,8 @@ def do_inference(
         prompter = Prompter(prompt_template)
         prompt = prompter.generate_prompt(variables)
 
-        if not lora_model_name:
-            lora_model_name = "None"
-        if "/" not in lora_model_name and lora_model_name != "None":
-            path_of_available_lora_model = get_path_of_available_lora_model(
-                lora_model_name)
-            if path_of_available_lora_model:
-                lora_model_name = path_of_available_lora_model
-
         if Global.ui_dev_mode:
-            message = f"Hi, I’m currently in UI-development mode and do not have access to resources to process your request. However, this behavior is similar to what will actually happen, so you can try and see how it will work!\n\nBase model: {Global.base_model}\nLoRA model: {lora_model_name}\n\nThe following text is your prompt:\n\n{prompt}"
+            message = f"Hi, I’m currently in UI-development mode and do not have access to resources to process your request. However, this behavior is similar to what will actually happen, so you can try and see how it will work!\n\nBase model: {base_model_name}\nLoRA model: {lora_model_name}\n\nThe following is your prompt:\n\n{prompt}"
             print(message)
 
             if stream_output:
@@ -90,18 +95,13 @@ def do_inference(
                 return
             time.sleep(1)
             yield (
-                gr.Textbox.update(value=message, lines=1), # TODO
+                gr.Textbox.update(value=message, lines=1),  # TODO
                 json.dumps(list(range(len(message.split()))), indent=2)
             )
             return
 
-        # model = get_base_model()
-        if lora_model_name != "None":
-            model = get_model_with_lora(lora_model_name)
-        else:
-            raise ValueError("No LoRA model selected.")
-
-        tokenizer = get_tokenizer()
+        tokenizer = get_tokenizer(base_model_name)
+        model = get_model(base_model_name, lora_model_name)
 
         inputs = tokenizer(prompt, return_tensors="pt")
         input_ids = inputs["input_ids"].to(device)
@@ -210,7 +210,6 @@ def do_inference(
             gr.Textbox.update(value=response, lines=inference_output_lines),
             raw_output)
 
-
     except Exception as e:
         raise gr.Error(e)
 
@@ -232,7 +231,7 @@ def reload_selections(current_lora_model, current_prompt_template):
 
     default_lora_models = ["tloen/alpaca-lora-7b"]
     available_lora_models = default_lora_models + get_available_lora_model_names()
-    available_lora_models = available_lora_models
+    available_lora_models = available_lora_models + ["None"]
 
     current_lora_model = current_lora_model or next(
         iter(available_lora_models), None)
@@ -462,6 +461,10 @@ def inference_ui():
         things_that_might_timeout.append(lora_model_change_event)
 
         generate_event = generate_btn.click(
+            fn=prepare_inference,
+            inputs=[lora_model],
+            outputs=[inference_output, inference_raw_output],
+        ).then(
             fn=do_inference,
             inputs=[
                 lora_model,

llama_lora/ui/main_page.py

@@ -1,7 +1,6 @@
 import gradio as gr
 
 from ..globals import Global
-from ..models import get_model_with_lora
 
 from .inference_ui import inference_ui
 from .finetune_ui import finetune_ui
@@ -31,7 +30,7 @@ def main_page():
             info = []
             if Global.version:
                 info.append(f"LLaMA-LoRA Tuner `{Global.version}`")
-            info.append(f"Base model: `{Global.base_model}`")
+            info.append(f"Base model: `{Global.default_base_model_name}`")
             if Global.ui_show_sys_info:
                 info.append(f"Data dir: `{Global.data_dir}`")
             gr.Markdown(f"""

llama_lora/ui/tokenizer_ui.py

@@ -7,11 +7,12 @@ from ..models import get_tokenizer
 
 
 def handle_decode(encoded_tokens_json):
+    base_model_name = Global.default_base_model_name
     try:
         encoded_tokens = json.loads(encoded_tokens_json)
         if Global.ui_dev_mode:
             return f"Not actually decoding tokens in UI dev mode.", gr.Markdown.update("", visible=False)
-        tokenizer = get_tokenizer()
+        tokenizer = get_tokenizer(base_model_name)
         decoded_tokens = tokenizer.decode(encoded_tokens)
         return decoded_tokens, gr.Markdown.update("", visible=False)
     except Exception as e:
@@ -19,10 +20,11 @@ def handle_decode(encoded_tokens_json):
 
 
 def handle_encode(decoded_tokens):
+    base_model_name = Global.default_base_model_name
     try:
         if Global.ui_dev_mode:
             return f"[\"Not actually encoding tokens in UI dev mode.\"]", gr.Markdown.update("", visible=False)
-        tokenizer = get_tokenizer()
+        tokenizer = get_tokenizer(base_model_name)
         result = tokenizer(decoded_tokens)
         encoded_tokens_json = json.dumps(result['input_ids'], indent=2)
         return encoded_tokens_json, gr.Markdown.update("", visible=False)

llama_lora/utils/lru_cache.py

@@ -25,3 +25,7 @@ class LRUCache:
 
     def clear(self):
         self.cache.clear()
+
+    def prepare_to_set(self):
+        if len(self.cache) >= self.capacity:
+            self.cache.popitem(last=False)