Merge pull request #32 from glide-the/quantized

load a 4-bit or 8-bit quantized version
THUDM · Oct 28, 2024 · 88c5c52 · 88c5c52
2 parents eb140fc + b5530ee
commit 88c5c52
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 2 deletions.
diff --git a/model_server.py b/model_server.py
@@ -7,7 +7,7 @@
 
 from fastapi import FastAPI, Request
 from fastapi.responses import StreamingResponse
-from transformers import AutoModel, AutoTokenizer
+from transformers import AutoModel, AutoTokenizer, BitsAndBytesConfig
 import torch
 import uvicorn
 
@@ -56,8 +56,16 @@ def __next__(self):
 class ModelWorker:
     def __init__(self, model_path, device='cuda'):
         self.device = device
+        bnb_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch.bfloat16
+        )
+
         self.glm_model = AutoModel.from_pretrained(model_path, trust_remote_code=True,
-                                                   device=device).to(device).eval()
+                                                   quantization_config=bnb_config,
+                                                   device_map={"":0}).eval()
         self.glm_tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
 
     @torch.inference_mode()

diff --git a/requirements.txt b/requirements.txt
@@ -34,3 +34,8 @@ wget==3.2
 WeTextProcessing==1.0.3
 torch==2.3.0
 torchaudio==2.3.0
+python-multipart==0.0.12
+
+#  load a 4-bit or 8-bit quantized version 
+bitsandbytes==0.44.1
+accelerate==1.0.1