.

app.py

@@ -19,8 +19,16 @@ if "torch_dtype" in model_config:
     elif model_config["torch_dtype"] == "bfloat16":
         model_config["torch_dtype"] = torch.bfloat16
 
+# Create quantization config
 quantization_config = BitsAndBytesConfig(load_in_8bit=True)
-model = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=quantization_config, **model_config)
+
+# Load the model with quantization config
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    quantization_config=quantization_config,
+    low_cpu_mem_usage=True,
+    **model_config
+)
 
 checkpoint_model = "checkpoint_dir/checkpoint-650"
 

config.yaml

@@ -5,4 +5,4 @@ model_config:
   use_cache: True
   attn_implementation: "eager"
   device_map: "cpu"
-  load_in_8bit: True
+  # Remove the load_in_8bit line