cache lora models

llama_lora/globals.py

@@ -6,6 +6,7 @@ from typing import Any, Dict, List, Optional, Tuple, Union
 from numba import cuda
 import nvidia_smi
 
+from .utils.lru_cache import LRUCache
 from .lib.finetune import train
 
 
@@ -31,8 +32,7 @@ class Global:
 
     # Model related
     model_has_been_used = False
-    loaded_base_model_with_lora = None
-    loaded_base_model_with_lora_name = None
+    cached_lora_models = LRUCache(10)
 
     # GPU Info
     gpu_cc = None  # GPU compute capability

llama_lora/models.py

@@ -31,30 +31,32 @@ def get_base_model():
     return Global.loaded_base_model
 
 
-def get_model_with_lora(lora_weights: str = "tloen/alpaca-lora-7b"):
+def get_model_with_lora(lora_weights_name_or_path: str = "tloen/alpaca-lora-7b"):
     Global.model_has_been_used = True
 
-    if Global.loaded_base_model_with_lora and Global.loaded_base_model_with_lora_name == lora_weights:
-        return Global.loaded_base_model_with_lora
+    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
 
     if device == "cuda":
         model = PeftModel.from_pretrained(
             get_base_model(),
-            lora_weights,
+            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_base_model(),
-            lora_weights,
+            lora_weights_name_or_path,
             device_map={"": device},
             torch_dtype=torch.float16,
         )
     else:
         model = PeftModel.from_pretrained(
             get_base_model(),
-            lora_weights,
+            lora_weights_name_or_path,
             device_map={"": device},
         )
 
@@ -69,8 +71,9 @@ def get_model_with_lora(lora_weights: str = "tloen/alpaca-lora-7b"):
     if torch.__version__ >= "2" and sys.platform != "win32":
         model = torch.compile(model)
 
-    Global.loaded_base_model_with_lora = model
-    Global.loaded_base_model_with_lora_name = lora_weights
+    if Global.cached_lora_models:
+        Global.cached_lora_models.set(lora_weights_name_or_path, model)
+
     return model
 
 
@@ -127,10 +130,7 @@ def unload_models():
     del Global.loaded_tokenizer
     Global.loaded_tokenizer = None
 
-    del Global.loaded_base_model_with_lora
-    Global.loaded_base_model_with_lora = None
-
-    Global.loaded_base_model_with_lora_name = None
+    Global.cached_lora_models.clear()
 
     clear_cache()
 

requirements.lock.txt

@@ -65,7 +65,7 @@ packaging==23.0
 pandas==2.0.0
 parso==0.8.3
 pathspec==0.11.1
-peft @ git+https://github.com/huggingface/peft.git@deff03f2c251534fffd2511fc2d440e84cc54b1b
+peft @ git+https://github.com/huggingface/peft.git@382b178911edff38c1ff619bbac2ba556bd2276b
 pexpect==4.8.0
 pickleshare==0.7.5
 Pillow==9.3.0