.

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, BitsAndBytesConfig
 
 from config import load_config
 
@@ -19,7 +19,8 @@ if "torch_dtype" in model_config:
     elif model_config["torch_dtype"] == "bfloat16":
         model_config["torch_dtype"] = torch.bfloat16
 
-model = AutoModelForCausalLM.from_pretrained(model_name, **model_config)
+quantization_config = BitsAndBytesConfig(load_in_8bit=True)
+model = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=quantization_config, **model_config)
 
 checkpoint_model = "checkpoint_dir/checkpoint-650"
 

config.yaml

@@ -4,4 +4,5 @@ model_config:
   trust_remote_code: True
   use_cache: True
   attn_implementation: "eager"
-  device_map: "cpu"
+  device_map: "cpu"
+  load_in_8bit: True