diff --git a/src/lighteval/models/transformers/transformers_model.py b/src/lighteval/models/transformers/transformers_model.py index 64e790a2f..c82a9689f 100644 --- a/src/lighteval/models/transformers/transformers_model.py +++ b/src/lighteval/models/transformers/transformers_model.py @@ -340,10 +340,13 @@ def init_model_parallel(self, model_parallel: bool | None = None) -> Tuple[bool, return False, None, None self.num_local_processes = int(os.environ.get("LOCAL_WORLD_SIZE", 1)) + # GPUs visible to each process, not the number of machines. One GPU per + # process is the usual data-parallel launch. An explicit model_parallel=True + # still shards the GPUs assigned to this process. self.num_machines = torch.cuda.device_count() // self.num_local_processes - if self.num_machines == 1: - logger.info("We are not in a distributed setting. Setting model_parallel to False.") + if self.num_machines == 1 and model_parallel is not True: + logger.info("Each process can see a single GPU, so model parallel stays off.") model_parallel = False if model_parallel is None: diff --git a/src/lighteval/models/transformers/vlm_transformers_model.py b/src/lighteval/models/transformers/vlm_transformers_model.py index 0697ab729..b4995a2d5 100644 --- a/src/lighteval/models/transformers/vlm_transformers_model.py +++ b/src/lighteval/models/transformers/vlm_transformers_model.py @@ -214,10 +214,13 @@ def init_model_parallel(self, model_parallel: bool | None = None) -> Tuple[bool, return False, None, None self.num_local_processes = int(os.environ.get("LOCAL_WORLD_SIZE", 1)) + # GPUs visible to each process, not the number of machines. One GPU per + # process is the usual data-parallel launch. An explicit model_parallel=True + # still shards the GPUs assigned to this process. self.num_machines = torch.cuda.device_count() // self.num_local_processes - if self.num_machines == 1: - logger.info("We are not in a distributed setting. Setting model_parallel to False.") + if self.num_machines == 1 and model_parallel is not True: + logger.info("Each process can see a single GPU, so model parallel stays off.") model_parallel = False if model_parallel is None: diff --git a/tests/unit/models/test_transformers_model.py b/tests/unit/models/test_transformers_model.py index da7c925ae..4d4532fd6 100644 --- a/tests/unit/models/test_transformers_model.py +++ b/tests/unit/models/test_transformers_model.py @@ -20,6 +20,7 @@ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. +import os import unittest from unittest.mock import Mock, patch @@ -426,5 +427,52 @@ 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 TestInitModelParallel(unittest.TestCase): + def _model(self, cls): + model = cls.__new__(cls) + accelerator = Mock() + accelerator.process_index = 0 + model.accelerator = accelerator + return model + + def _run(self, cls, module: str, local_world_size: str, device_count: int, model_parallel): + model = self._model(cls) + memory = {0: "10GiB", 1: "10GiB"} + with ( + patch.dict(os.environ, {"LOCAL_WORLD_SIZE": local_world_size}), + patch(f"{module}.torch.cuda.device_count", return_value=device_count), + patch(f"{module}.get_max_memory", return_value=dict(memory)), + ): + return model.init_model_parallel(model_parallel) + + def test_explicit_model_parallel_is_kept_when_each_process_has_one_gpu(self): + from lighteval.models.transformers.transformers_model import TransformersModel + from lighteval.models.transformers.vlm_transformers_model import VLMTransformersModel + + cases = [ + (TransformersModel, "lighteval.models.transformers.transformers_model"), + (VLMTransformersModel, "lighteval.models.transformers.vlm_transformers_model"), + ] + for cls, module in cases: + with self.subTest(cls=cls.__name__): + enabled, max_memory, device_map = self._run(cls, module, "2", 2, True) + self.assertTrue(enabled) + self.assertEqual(device_map, "auto") + self.assertEqual(max_memory, {0: "10GiB"}) + + enabled, max_memory, device_map = self._run(cls, module, "2", 2, None) + self.assertFalse(enabled) + self.assertIsNone(max_memory) + self.assertIsNone(device_map) + + def test_unset_model_parallel_uses_every_gpu_from_one_process(self): + enabled, max_memory, device_map = self._run( + TransformersModel, "lighteval.models.transformers.transformers_model", "1", 2, None + ) + self.assertTrue(enabled) + self.assertEqual(device_map, "auto") + self.assertEqual(max_memory, {0: "10GiB", 1: "10GiB"}) + + if __name__ == "__main__": unittest.main()