From 600bd7841d7f24635b3472b7ec56fa2b22c69870 Mon Sep 17 00:00:00 2001 From: Ella Charlaix Date: Mon, 4 Mar 2024 15:20:52 +0100 Subject: [PATCH 1/2] Add optimum models support --- run_evals_accelerate.py | 7 +++++- src/lighteval/models/model_config.py | 13 +++++++++++ src/lighteval/models/model_loader.py | 17 +++++++++++--- src/lighteval/models/optimum_model.py | 33 +++++++++++++++++++++++++++ src/lighteval/utils.py | 3 +++ 5 files changed, 69 insertions(+), 4 deletions(-) create mode 100644 src/lighteval/models/optimum_model.py diff --git a/run_evals_accelerate.py b/run_evals_accelerate.py index 474936be..11cf60a2 100644 --- a/run_evals_accelerate.py +++ b/run_evals_accelerate.py @@ -26,10 +26,15 @@ def get_parser(): default=False, help="set to True of your model has been trained with peft, also need to provide the base model name", ) + weight_type_group.add_argument( + "--optimum_weights", + action="store_true", + default=False, + help="set to True of your model needs to be loadded with optimum", + ) parser.add_argument( "--base_model", type=str, default=None, help="name of the base model to be used for delta or adapter weights" ) - task_type_group.add_argument("--model_args") parser.add_argument("--model_dtype", type=str, default=None) parser.add_argument( diff --git a/src/lighteval/models/model_config.py b/src/lighteval/models/model_config.py index f1bbef43..0027cb89 100644 --- a/src/lighteval/models/model_config.py +++ b/src/lighteval/models/model_config.py @@ -10,10 +10,12 @@ from lighteval.utils import ( NO_AUTOGPTQ_ERROR_MSG, NO_BNB_ERROR_MSG, + NO_OPTIMUM_ERROR_MSG, NO_PEFT_ERROR_MSG, is_accelerate_available, is_autogptq_available, is_bnb_available, + is_optimum_available, is_peft_available, ) @@ -190,6 +192,15 @@ def init_configs(self, env_config: EnvConfig): return self._init_configs(self.base_model, env_config) +@dataclass +class OptimumModelConfig(BaseModelConfig): + def __post_init__(self): + if not is_optimum_available(): + raise ImportError(NO_OPTIMUM_ERROR_MSG) + + return super().__post_init__() + + @dataclass class TGIModelConfig: inference_server_address: str @@ -298,6 +309,8 @@ def create_model_config(args: Namespace, accelerator: Union["Accelerator", None] if args.base_model is None: raise ValueError("You need to specify a base model when using adapter weights") return AdapterModelConfig(**args_dict) + if args.optimum_weights: + return OptimumModelConfig(**args_dict) if args.base_model is not None: raise ValueError("You can't specifify a base model if you are not using delta/adapter weights") return BaseModelConfig(**args_dict) diff --git a/src/lighteval/models/model_loader.py b/src/lighteval/models/model_loader.py index 5bb7d091..273154ee 100644 --- a/src/lighteval/models/model_loader.py +++ b/src/lighteval/models/model_loader.py @@ -13,8 +13,10 @@ EnvConfig, InferenceEndpointModelConfig, InferenceModelConfig, + OptimumModelConfig, TGIModelConfig, ) +from lighteval.models.optimum_model import OptimumModel from lighteval.models.tgi_model import ModelClient from lighteval.utils import NO_TGI_ERROR_MSG, is_accelerate_available, is_tgi_available @@ -32,9 +34,16 @@ class ModelInfo: def load_model( # noqa: C901 - config: Union[BaseModelConfig, AdapterModelConfig, DeltaModelConfig, TGIModelConfig, InferenceEndpointModelConfig], + config: Union[ + BaseModelConfig, + AdapterModelConfig, + DeltaModelConfig, + TGIModelConfig, + InferenceEndpointModelConfig, + OptimumModelConfig, + ], env_config: EnvConfig, -) -> Tuple[Union[BaseModel, AdapterModel, DeltaModel, ModelClient], ModelInfo]: +) -> Tuple[Union[BaseModel, AdapterModel, DeltaModel, ModelClient, OptimumModel], ModelInfo]: """Will load either a model from an inference server or a model from a checkpoint. depending on the arguments passed to the program. @@ -93,12 +102,14 @@ def load_model_with_inference_endpoints(config: InferenceEndpointModelConfig, en def load_model_with_accelerate_or_default( - config: Union[AdapterModelConfig, BaseModelConfig, DeltaModelConfig], env_config: EnvConfig + config: Union[AdapterModelConfig, BaseModelConfig, DeltaModelConfig, OptimumModelConfig], env_config: EnvConfig ): if isinstance(config, AdapterModelConfig): model = AdapterModel(config=config, env_config=env_config) elif isinstance(config, DeltaModelConfig): model = DeltaModel(config=config, env_config=env_config) + elif isinstance(config, OptimumModelConfig): + model = OptimumModel(config=config, env_config=env_config) else: model = BaseModel(config=config, env_config=env_config) diff --git a/src/lighteval/models/optimum_model.py b/src/lighteval/models/optimum_model.py new file mode 100644 index 00000000..3d7a943d --- /dev/null +++ b/src/lighteval/models/optimum_model.py @@ -0,0 +1,33 @@ +from lighteval.models.base_model import BaseModel +from lighteval.models.model_config import EnvConfig, OptimumModelConfig +from lighteval.models.utils import _get_dtype +from lighteval.utils import is_optimum_available + + +if is_optimum_available(): + # from optimum import OptimumConfig + from optimum.intel.openvino import OVModelForCausalLM + + +class OptimumModel(BaseModel): + + + def _create_auto_model(self, config: OptimumModelConfig, env_config: EnvConfig): + + # TODO : Get loading class from optimum config (add support for ORTModelForCausalLM / INCModelForCausalLM / IPEXModelForCausalLM) + # optimum_config = OptimumConfig.from_pretrained(config.pretrained) + + config.model_parallel, max_memory, device_map = self.init_model_parallel(config.model_parallel) + torch_dtype = _get_dtype(config.dtype, self._config) + + model = OVModelForCausalLM.from_pretrained( + config.pretrained, + revision=config.revision + (f"/{config.subfolder}" if config.subfolder is not None else ""), + torch_dtype=torch_dtype, + trust_remote_code=config.trust_remote_code, + cache_dir=env_config.cache_dir, + use_auth_token=env_config.token, + quantization_config=config.quantization_config, + ) + + return model diff --git a/src/lighteval/utils.py b/src/lighteval/utils.py index c2a9335d..bdb619e1 100644 --- a/src/lighteval/utils.py +++ b/src/lighteval/utils.py @@ -170,6 +170,9 @@ def is_optimum_available() -> bool: return importlib.util.find_spec("optimum") is not None +NO_OPTIMUM_ERROR_MSG = "You are trying to load a model with `optimum`, which is not available in your local environement. Please install it using pip." + + def is_bnb_available() -> bool: return importlib.util.find_spec("bitsandbytes") is not None From 9536181cee0cd51afb1a0a9f0537a8c124fa3a06 Mon Sep 17 00:00:00 2001 From: Ella Charlaix Date: Wed, 6 Mar 2024 18:01:28 +0100 Subject: [PATCH 2/2] fix format --- src/lighteval/models/optimum_model.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/src/lighteval/models/optimum_model.py b/src/lighteval/models/optimum_model.py index 3d7a943d..ee734fa8 100644 --- a/src/lighteval/models/optimum_model.py +++ b/src/lighteval/models/optimum_model.py @@ -10,10 +10,7 @@ class OptimumModel(BaseModel): - - def _create_auto_model(self, config: OptimumModelConfig, env_config: EnvConfig): - # TODO : Get loading class from optimum config (add support for ORTModelForCausalLM / INCModelForCausalLM / IPEXModelForCausalLM) # optimum_config = OptimumConfig.from_pretrained(config.pretrained)