Skip batch size search when model parallel has several processes

src/lighteval/models/transformers/transformers_model.py

@@ -505,6 +505,15 @@ def _check_continuations_start_space(self, continuation: str) -> str:
     def _get_batch_size(self, max_input_length: int, override_bs: int | None, starting_batch_size: int = 512) -> int:
         if override_bs is not None:
             return override_bs
+        # The search runs a forward and halves the batch after an OOM. Ranks do not
+        # take the same path, so under model parallel that search deadlocks.
+        num_processes = getattr(self.accelerator, "num_processes", 1)
+        if self.config.model_parallel and num_processes > 1:
+            logger.info(
+                "Model parallel is on for %s processes, so batch size detection is skipped. Using batch size 1.",
+                num_processes,
+            )
+            return 1
         logger.info(f"Detecting largest batch size with max_input_length={max_input_length}")
 
         @find_executable_batch_size(

tests/unit/models/test_transformers_model.py

@@ -426,5 +426,62 @@ def test_transformers_model_use_chat_template_with_different_model_names(
                 self.assertEqual(model.use_chat_template, model._tokenizer.chat_template is not None)
 
 
+class TestGetBatchSize(unittest.TestCase):
+    def _model(self, model_parallel, num_processes):
+        model = TransformersModel.__new__(TransformersModel)
+        model.config = TransformersModelConfig(model_name="gpt2", model_parallel=model_parallel)
+        model.accelerator = Mock()
+        model.accelerator.num_processes = num_processes
+        return model
+
+    def test_distributed_model_parallel_does_not_search_for_a_batch_size(self):
+        model = self._model(model_parallel=True, num_processes=2)
+        with patch("lighteval.models.transformers.transformers_model.find_executable_batch_size") as search:
+            batch_size = model._get_batch_size(max_input_length=32, override_bs=None)
+        self.assertEqual(batch_size, 1)
+        search.assert_not_called()
+
+    def test_an_explicit_batch_size_is_kept_under_model_parallel(self):
+        model = self._model(model_parallel=True, num_processes=2)
+        with patch("lighteval.models.transformers.transformers_model.find_executable_batch_size") as search:
+            batch_size = model._get_batch_size(max_input_length=32, override_bs=8)
+        self.assertEqual(batch_size, 8)
+        search.assert_not_called()
+
+    def test_data_parallel_still_searches_for_a_batch_size(self):
+        model = self._model(model_parallel=False, num_processes=2)
+
+        def decorator(*args, **kwargs):
+            def wrap(function):
+                return lambda: 16
+
+            return wrap
+
+        with patch(
+            "lighteval.models.transformers.transformers_model.find_executable_batch_size",
+            side_effect=decorator,
+        ) as search:
+            batch_size = model._get_batch_size(max_input_length=32, override_bs=None)
+        self.assertEqual(batch_size, 16)
+        search.assert_called_once()
+
+    def test_one_process_still_searches_for_a_batch_size(self):
+        model = self._model(model_parallel=True, num_processes=1)
+
+        def decorator(*args, **kwargs):
+            def wrap(function):
+                return lambda: 32
+
+            return wrap
+
+        with patch(
+            "lighteval.models.transformers.transformers_model.find_executable_batch_size",
+            side_effect=decorator,
+        ) as search:
+            batch_size = model._get_batch_size(max_input_length=32, override_bs=None)
+        self.assertEqual(batch_size, 32)
+        search.assert_called_once()
+
+
 if __name__ == "__main__":
     unittest.main()