fix slow tests (#689)

tests on the aws runner were hanging, culprit was multiporcessing when loading datasets.

Author: NathanHB

.github/workflows/slow_tests.yaml

@@ -37,11 +37,6 @@ jobs:
       - name: Install the project
         run: uv sync --extra dev
 
-      - name: Ensure cache directories exist
-        run: mkdir -p cache/models cache/datasets
 
       - name: Run tests
-        env:
-          HF_HOME: "cache/models"
-          HF_DATASETS_CACHE: "cache/datasets"
-        run: uv run pytest --disable-pytest-warnings --runslow tests/slow_tests
+        run: uv run pytest --disable-pytest-warnings -o log_cli=true -o log_cli_level=INFO --runslow tests/slow_tests/

examples/model_configs/transformers_model.yaml

@@ -4,6 +4,7 @@ model_parameters:
   dtype: "float16"
   compile: false
   model_parallel: false
+  batch_size: 1
   multichoice_continuations_start_space: null # If true/false, will force multiple choice continuations to start/not start with a space. If none, will do nothing
   generation_parameters:
     temperature: 0.2

src/lighteval/models/transformers/transformers_model.py

@@ -226,6 +226,12 @@ def __init__(
             model_size=str(model_size),
         )
 
+    def cleanup(self):
+        """Clean up operations if needed, such as closing an endpoint."""
+        del self.model
+        del self._tokenizer
+        torch.cuda.empty_cache()
+
     @classmethod
     def from_model(
         cls,
@@ -543,7 +549,7 @@ def greedy_until(
                     longest_context_continuation_size_in_split, self.max_length
                 )
             batch_size = self._get_batch_size(
-                override_bs=self.batch_size,
+                override_bs=self.config.batch_size,
                 max_input_length=max_context_continuation_size_allowed,
                 starting_batch_size=starting_batch_size,
             )
@@ -710,7 +716,6 @@ def _generate(
     def loglikelihood(
         self,
         requests: list[LoglikelihoodRequest],
-        override_bs: Optional[int] = None,
     ) -> list[LoglikelihoodResponse]:
         """Tokenize the context and continuation and compute the log likelihood of those
         tokenized sequences.
@@ -731,12 +736,11 @@ def loglikelihood(
                     request.context, request.choice, pairwise=self.pairwise_tokenization
                 )
 
-        return self._loglikelihood_tokens(requests, override_bs=override_bs)
+        return self._loglikelihood_tokens(requests)
 
     def loglikelihood_rolling(
         self,
         requests: list[LoglikelihoodRollingRequest],
-        override_bs=None,
     ) -> list[LoglikelihoodResponse]:
         """This function is used to compute the log likelihood of the context for perplexity metrics."""
 
@@ -746,7 +750,6 @@ def loglikelihood_rolling(
 
         results = self._loglikelihood_tokens(
             requests,
-            override_bs=override_bs,
             return_bool_score=False,
             rolling=True,
         )
@@ -755,7 +758,6 @@ def loglikelihood_rolling(
     def _loglikelihood_tokens(
         self,
         requests: list[LoglikelihoodRequest],
-        override_bs: int = -1,
         return_bool_score: bool = True,
         rolling: bool = False,
     ) -> list[LoglikelihoodResponse]:
@@ -774,7 +776,7 @@ def _loglikelihood_tokens(
                 )
 
             batch_size = self._get_batch_size(
-                override_bs=override_bs,
+                override_bs=self.config.batch_size,
                 max_input_length=max_context_continuation_size_allowed,
                 starting_batch_size=starting_batch_size,
             )
@@ -967,7 +969,8 @@ def pad_and_gather(
         return output_tensor, length_tensor
 
     def loglikelihood_single_token(
-        self, requests: list[LoglikelihoodSingleTokenRequest], override_bs: Optional[int] = None
+        self,
+        requests: list[LoglikelihoodSingleTokenRequest],
     ) -> list[LoglikelihoodSingleTokenResponse]:
         """Tokenize the context and continuation and compute the log likelihood of those
         tokenized sequences.
@@ -996,10 +999,11 @@ def loglikelihood_single_token(
                 )
             request.tokenized_continuation = continuations_enc
 
-        return self._loglikelihood_single_token(requests, override_bs=override_bs)
+        return self._loglikelihood_single_token(requests)
 
     def _loglikelihood_single_token(
-        self, requests: list[LoglikelihoodSingleTokenRequest], override_bs: int = -1
+        self,
+        requests: list[LoglikelihoodSingleTokenRequest],
     ) -> list[LoglikelihoodSingleTokenResponse]:
         dataset = LoglikelihoodSingleTokenDataset(requests=requests, num_dataset_splits=self.DATASET_SPLITS)
         starting_batch_size = STARTING_BATCH_SIZE
@@ -1008,7 +1012,7 @@ def _loglikelihood_single_token(
         for split_start, split_end in tqdm(dataset.splits_start_end_iterator()):
             context_enc = dataset[0].tokenized_context
             max_context = len(context_enc[-self.max_length :])
-            batch_size = self._get_batch_size(override_bs=override_bs, max_input_length=max_context)
+            batch_size = self._get_batch_size(override_bs=self.config.batch_size, max_input_length=max_context)
             starting_batch_size = batch_size * 2
 
             dataloader = DataLoader(dataset, batch_size=starting_batch_size, collate_fn=lambda batch: batch)

tests/slow_tests/test_accelerate_model.py

@@ -56,7 +56,7 @@ def run_model(model_name: str, use_chat_template: bool):
         tasks=TASKS_PATH,
         use_chat_template=use_chat_template,
         output_dir="",
-        dataset_loading_processes=8,
+        dataset_loading_processes=1,
         save_details=False,
         max_samples=10,
         custom_tasks=CUSTOM_TASKS_PATH,

tests/slow_tests/test_vllm_model.py

@@ -56,7 +56,7 @@ def run_model(model_name: str, use_chat_template: bool):
         tasks=TASKS_PATH,
         use_chat_template=use_chat_template,
         output_dir="",
-        dataset_loading_processes=8,
+        dataset_loading_processes=1,
         save_details=False,
         max_samples=10,
         custom_tasks=CUSTOM_TASKS_PATH,