preload base model on colab

LLaMA_LoRA.ipynb

@@ -116,7 +116,7 @@
         "# @markdown You can customize the location of the stored data here.\n",
         "\n",
         "# @markdown The folder in Google Drive where Colab Notebook data are stored<br />&nbsp;&nbsp;**(WARNING: The content of this folder will be modified by this notebook)**:\n",
-        "google_drive_folder = \"Colab Data/LLaMA LoRA\" # @param {type:\"string\"}\n",
+        "google_drive_folder = \"Colab Data/LLaMA-LoRA Tuner\" # @param {type:\"string\"}\n",
         "# google_drive_colab_data_folder = \"Colab Notebooks/Notebook Data\"\n",
         "\n",
         "# Where Google Drive will be mounted in the Colab runtime.\n",
@@ -289,7 +289,11 @@
         "# Prepare Data Dir\n",
         "import os\n",
         "from llama_lora.llama_lora.utils.data import init_data_dir\n",
-        "init_data_dir()"
+        "init_data_dir()",
+        "\n",
+        "# Load the Base Model\n",
+        "from llama_lora.llama_lora.models import prepare_base_model\n",
+        "prepare_base_model()\n"
       ],
       "metadata": {
         "id": "Yf6g248ylteP"

llama_lora/globals.py

@@ -31,6 +31,8 @@ class Global:
     # Model related
     loaded_models = LRUCache(1)
     loaded_tokenizers = LRUCache(1)
+    new_base_model_that_is_ready_to_be_used = None
+    name_of_new_base_model_that_is_ready_to_be_used = None
 
     # GPU Info
     gpu_cc = None  # GPU compute capability

llama_lora/models.py

@@ -26,6 +26,17 @@ def get_new_base_model(base_model_name):
     if Global.ui_dev_mode:
         return
 
+    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:
+            model = Global.new_base_model_that_is_ready_to_be_used
+            Global.new_base_model_that_is_ready_to_be_used = None
+            Global.name_of_new_base_model_that_is_ready_to_be_used = None
+            return model
+        else:
+            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()
+
     device = get_device()
 
     if device == "cuda":
@@ -34,7 +45,8 @@ def get_new_base_model(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
+            # ? https://github.com/tloen/alpaca-lora/issues/21
+            device_map={'': 0},
         )
     elif device == "mps":
         model = LlamaForCausalLM.from_pretrained(
@@ -69,8 +81,8 @@ def get_tokenizer(base_model_name):
 
 
 def get_model(
-    base_model_name,
-    peft_model_name = None):
+        base_model_name,
+        peft_model_name=None):
     if Global.ui_dev_mode:
         return
 
@@ -88,7 +100,8 @@ def get_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)
+    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
 
@@ -105,7 +118,8 @@ def get_model(
                 model,
                 peft_model_name_or_path,
                 torch_dtype=torch.float16,
-                device_map={'': 0},  # ? https://github.com/tloen/alpaca-lora/issues/21
+                # ? https://github.com/tloen/alpaca-lora/issues/21
+                device_map={'': 0},
             )
         elif device == "mps":
             model = PeftModel.from_pretrained(
@@ -138,6 +152,11 @@ def get_model(
     return model
 
 
+def prepare_base_model(base_model_name=Global.default_base_model_name):
+    Global.new_base_model_that_is_ready_to_be_used = get_new_base_model(base_model_name)
+    Global.name_of_new_base_model_that_is_ready_to_be_used = base_model_name
+
+
 def clear_cache():
     gc.collect()