Allow loading model with 4bit quantization.

For detail on 4bit options, see: https://huggingface.co/blog/4bit-transformers-bitsandbytes

Author: idoru

basaran/__init__.py

@@ -21,6 +21,9 @@ def is_true(value):
 MODEL_REVISION = os.getenv("MODEL_REVISION", "")
 MODEL_CACHE_DIR = os.getenv("MODEL_CACHE_DIR", "models")
 MODEL_LOAD_IN_8BIT = is_true(os.getenv("MODEL_LOAD_IN_8BIT", ""))
+MODEL_LOAD_IN_4BIT = is_true(os.getenv("MODEL_LOAD_IN_4BIT", ""))
+MODEL_4BIT_QUANT_TYPE = os.getenv("MODEL_4BIT_QUANT_TYPE", "fp4")
+MODEL_4BIT_DOUBLE_QUANT = is_true(os.getenv("MODEL_4BIT_DOUBLE_QUANT", ""))
 MODEL_LOCAL_FILES_ONLY = is_true(os.getenv("MODEL_LOCAL_FILES_ONLY", ""))
 MODEL_TRUST_REMOTE_CODE = is_true(os.getenv("MODEL_TRUST_REMOTE_CODE", ""))
 MODEL_HALF_PRECISION = is_true(os.getenv("MODEL_HALF_PRECISION", ""))

basaran/__main__.py

@@ -20,6 +20,9 @@
 from . import MODEL_REVISION
 from . import MODEL_CACHE_DIR
 from . import MODEL_LOAD_IN_8BIT
+from . import MODEL_LOAD_IN_4BIT
+from . import MODEL_4BIT_QUANT_TYPE
+from . import MODEL_4BIT_DOUBLE_QUANT
 from . import MODEL_LOCAL_FILES_ONLY
 from . import MODEL_TRUST_REMOTE_CODE
 from . import MODEL_HALF_PRECISION
@@ -42,6 +45,9 @@
     revision=MODEL_REVISION,
     cache_dir=MODEL_CACHE_DIR,
     load_in_8bit=MODEL_LOAD_IN_8BIT,
+    load_in_4bit=MODEL_LOAD_IN_4BIT,
+    quant_type=MODEL_4BIT_QUANT_TYPE,
+    double_quant=MODEL_4BIT_DOUBLE_QUANT,
     local_files_only=MODEL_LOCAL_FILES_ONLY,
     trust_remote_code=MODEL_TRUST_REMOTE_CODE,
     half_precision=MODEL_HALF_PRECISION,

basaran/model.py

@@ -12,6 +12,7 @@
     MinNewTokensLengthLogitsProcessor,
     TemperatureLogitsWarper,
     TopPLogitsWarper,
+    BitsAndBytesConfig
 )
 
 from .choice import map_choice
@@ -302,6 +303,9 @@ def load_model(
     revision=None,
     cache_dir=None,
     load_in_8bit=False,
+    load_in_4bit=False,
+    quant_type="fp4",
+    double_quant=False,
     local_files_only=False,
     trust_remote_code=False,
     half_precision=False,
@@ -319,12 +323,27 @@ def load_model(
 
     # Set device mapping and quantization options if CUDA is available.
     if torch.cuda.is_available():
+        # Set quantization options if specified.
+        quant_config = None
+        if load_in_8bit and load_in_4bit:
+            raise ValueError("Only one of load_in_8bit and load_in_4bit can be True")
+        if load_in_8bit:
+            quant_config = BitsAndBytesConfig(
+                load_in_8bit=True,
+            )
+        elif load_in_4bit:
+            quant_config = BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_quant_type=quant_type,
+                bnb_4bit_use_double_quant=double_quant,
+                bnb_4bit_compute_dtype=torch.bfloat16,
+            )
         kwargs = kwargs.copy()
         kwargs["device_map"] = "auto"
-        kwargs["load_in_8bit"] = load_in_8bit
+        kwargs["quantization_config"] = quant_config
 
         # Cast all parameters to float16 if quantization is enabled.
-        if half_precision or load_in_8bit:
+        if half_precision or load_in_8bit or load_in_4bit:
             kwargs["torch_dtype"] = torch.float16
 
     # Support both decoder-only and encoder-decoder models.