Keep explicit model parallel on one multi-GPU machine

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:

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:

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()