Uses torch compile if provided

examples/model_configs/base_model.yaml

@@ -3,6 +3,7 @@ model:
   base_params:
     model_args: "pretrained=HuggingFaceH4/zephyr-7b-beta,revision=main" # pretrained=model_name,trust_remote_code=boolean,revision=revision_to_use,model_parallel=True ...
     dtype: "bfloat16"
+    compile: true
   merged_weights: # Ignore this section if you are not using PEFT models
     delta_weights: false # set to True of your model should be merged with a base model, also need to provide the base model name
     adapter_weights: false # set to True of your model has been trained with peft, also need to provide the base model name

src/lighteval/models/base_model.py

@@ -91,6 +91,8 @@ def __init__(
         if not config.model_parallel and not isinstance(config.quantization_config, BitsAndBytesConfig):
             hlog(f"Using Data Parallelism, putting model on device {self._device}")
             self.model = self.model.to(self._device)
+        if config.compile:
+            self.model.model.compile()
 
         self.model_name = _simplify_name(config.pretrained)
         self.model_sha = config.get_model_sha()

src/lighteval/models/model_config.py

@@ -124,6 +124,7 @@ class BaseModelConfig:
     quantization_config: Optional[BitsAndBytesConfig] = None
     trust_remote_code: bool = False
     use_chat_template: bool = False
+    compile: bool = False
 
     def __post_init__(self):
         if self.quantization_config is not None and not is_bnb_available():
@@ -309,6 +310,7 @@ def create_model_config(  # noqa: C901
 
         args_dict["accelerator"] = accelerator
         args_dict["use_chat_template"] = args.use_chat_template
+        args_dict["compile"] = bool(args_dict["compile"]) if "compile" in args_dict else False
 
         return BaseModelConfig(**args_dict)
 
@@ -371,6 +373,7 @@ def create_model_config(  # noqa: C901
 
         # We store the relevant other args
         args_dict["base_model"] = config["merged_weights"]["base_model"]
+        args_dict["compile"] = bool(config["base_params"]["compile"])
         args_dict["dtype"] = config["base_params"]["dtype"]
         args_dict["accelerator"] = accelerator
         args_dict["quantization_config"] = quantization_config