feat: Add low cpu by default (#40)

Signed-off-by: Ce Gao <cegao@tensorchord.ai>
tensorchord · May 25, 2023 · 59e2524 · 59e2524
1 parent 0f68f8a
commit 59e2524
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/src/modelz_llm/falcon_service.py b/src/modelz_llm/falcon_service.py
@@ -39,10 +39,12 @@ class LLM:
     def __init__(self, model_name: str, device: str) -> None:
         self.model_name = model_name
         self.tokenizer = transformers.AutoTokenizer.from_pretrained(
-            model_name, trust_remote_code=True
+            model_name, trust_remote_code=True, low_cpu_mem_usage=True
         )
         model_cls = getattr(transformers, LanguageModels.transformer_cls(model_name))
-        self.model = model_cls.from_pretrained(model_name, trust_remote_code=True)
+        self.model = model_cls.from_pretrained(
+            model_name, trust_remote_code=True, low_cpu_mem_usage=True
+        )
         if device == "auto":
             self.device = (
                 torch.cuda.current_device() if torch.cuda.is_available() else "cpu"