From c71ac979a19619feb47e7f6d366a05700f004585 Mon Sep 17 00:00:00 2001 From: Chetana Tanmayee Tadikamalla Date: Fri, 23 Oct 2026 10:53:50 +0530 Subject: [PATCH] Skip batch size search when model parallel has several processes The search runs a forward and halves the batch after an out-of-memory error. Ranks do not take the same path, so the search deadlocks. Fixes #896 Signed-off-by: Chetana Tanmayee Tadikamalla --- .../models/transformers/transformers_model.py | 9 +++ tests/unit/models/test_transformers_model.py | 57 +++++++++++++++++++ 2 files changed, 66 insertions(+) diff --git a/src/lighteval/models/transformers/transformers_model.py b/src/lighteval/models/transformers/transformers_model.py index 64e790a2f..834cee268 100644 --- a/src/lighteval/models/transformers/transformers_model.py +++ b/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( diff --git a/tests/unit/models/test_transformers_model.py b/tests/unit/models/test_transformers_model.py index da7c925ae..6f8f1b1ca 100644 --- a/tests/unit/models/test_transformers_model.py +++ b/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()