Added support for launching inference endpoint with different model dtypes (#124)

* Added support for any dtype

---------

Co-authored-by: Clémentine Fourrier <[email protected]>

Author: shaltielshmid

src/lighteval/models/endpoint_model.py

@@ -87,6 +87,7 @@ def __init__(
                             "MAX_INPUT_LENGTH": "2047",
                             "MAX_TOTAL_TOKENS": "2048",
                             "MODEL_ID": "/repository",
+                            **config.get_dtype_args(),
                         },
                         "url": "ghcr.io/huggingface/text-generation-inference:1.1.0",
                     },

src/lighteval/models/model_config.py

@@ -22,7 +22,7 @@
 
 from argparse import Namespace
 from dataclasses import dataclass
-from typing import Optional, Union
+from typing import Dict, Optional, Union
 
 import torch
 from transformers import AutoConfig, BitsAndBytesConfig, GPTQConfig, PretrainedConfig
@@ -233,11 +233,24 @@ class InferenceEndpointModelConfig:
     region: str
     instance_size: str
     instance_type: str
+    model_dtype: str
     framework: str = "pytorch"
     endpoint_type: str = "protected"
     should_reuse_existing: bool = False
     add_special_tokens: bool = True
 
+    def get_dtype_args(self) -> Dict[str, str]:
+        model_dtype = self.model_dtype.lower()
+        if model_dtype in ["awq", "eetq", "gptq"]:
+            return {"QUANTIZE": model_dtype}
+        if model_dtype == "8bit":
+            return {"QUANTIZE": "bitsandbytes"}
+        if model_dtype == "4bit":
+            return {"QUANTIZE": "bitsandbytes-nf4"}
+        if model_dtype in ["bfloat16", "float16"]:
+            return {"DTYPE": model_dtype}
+        return {}
+
 
 def create_model_config(args: Namespace, accelerator: Union["Accelerator", None]) -> BaseModelConfig:  # noqa: C901
     """
@@ -282,6 +295,7 @@ def create_model_config(args: Namespace, accelerator: Union["Accelerator", None]
                 instance_size=args.instance_size,
                 instance_type=args.instance_type,
                 should_reuse_existing=args.reuse_existing,
+                model_dtype=args.model_dtype,
             )
         return InferenceModelConfig(model=args.endpoint_model_name)
 

src/lighteval/models/model_loader.py

@@ -108,7 +108,7 @@ def load_model_with_inference_endpoints(config: InferenceEndpointModelConfig, en
     model_info = ModelInfo(
         model_name=model.name,
         model_sha=model.revision,
-        model_dtype="default",
+        model_dtype=config.model_dtype or "default",
         model_size=-1,
     )
     return model, model_info