Set local_devices from CUDA_VISIBLE_DEVICES when initializing distributed processing on GPUs

Author: gabeweisz

src/maxtext/utils/max_utils.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-""" Common Max Utils needed by multiple modules.
+"""Common Max Utils needed by multiple modules.
 All the functions include MaxText modules, such as Pyconfig, should be moved to MaxText utils file."""
 
 import collections
@@ -248,11 +248,20 @@ def initialize_jax_for_gpu(raw_keys):
   if os.environ.get("JAX_COORDINATOR_IP") is not None:
     coordinator_ip = str(os.getenv("JAX_COORDINATOR_IP"))
     coordinator_port = str(os.getenv("JAX_COORDINATOR_PORT"))
+    devices = os.getenv("CUDA_VISIBLE_DEVICES")
+    if devices is not None:
+      try:
+        devices = [int(x) for x in devices.split(",")]
+      except (ValueError, TypeError) as e:
+        max_logging.log(f"Error parsing CUDA_VISIBLE_DEVICES: {e}")
+        devices = None
+
     jax.distributed.initialize(
         coordinator_address=f"{coordinator_ip}:{coordinator_port}",
         num_processes=int(os.getenv("NNODES")),
         process_id=int(os.getenv("NODE_RANK")),
         initialization_timeout=raw_keys["jax_distributed_initialization_timeout"],
+        local_device_ids=devices,
     )
     max_logging.log(f"JAX global devices: {jax.devices()}")
 

tests/unit/max_utils_test.py

@@ -12,12 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-""" Tests for the common Max Utils """
+"""Tests for the common Max Utils"""
+import os
 import sys
 import unittest
 import time
 import pytest
 
+
 import jax
 from jax import numpy as jnp
 from jax import random
@@ -30,6 +32,7 @@
 from maxtext.utils import max_utils
 from maxtext.utils.train_utils import setup_train_loop
 from tests.utils.test_helpers import get_test_config_path
+from unittest import mock
 
 
 class MaxUtilsSummaryStats(unittest.TestCase):
@@ -168,5 +171,73 @@ def test_unscan_train_state_params(self):
     self.assertNotIn("layers_0", state.params["params"]["decoder"])
 
 
+class TestGpuDistributedInitialization(unittest.TestCase):
+  """Tests using CUDA_VISIBLE_DEVICES to control which GPUs are used in jax.distributed.initialize."""
+
+  @mock.patch.dict(
+      os.environ,
+      {
+          "JAX_COORDINATOR_IP": "10.0.0.1",
+          "JAX_COORDINATOR_PORT": "1234",
+          "NNODES": "1",
+          "NODE_RANK": "0",
+          "CUDA_VISIBLE_DEVICES": "0,2,3",  # Simulating Slurm/orchestrator assignment
+      },
+  )
+  @mock.patch("jax.distributed.initialize")
+  @mock.patch("jax.devices")
+  @mock.patch("maxtext.utils.max_logging.log")
+  def test_initialize_jax_for_gpu_valid_devices(self, _mock_log, _mock_devices, mock_init):
+    """Verifies that a comma-separated string of IDs is correctly parsed."""
+    raw_keys = {"jax_distributed_initialization_timeout": 300}
+    max_utils.initialize_jax_for_gpu(raw_keys)
+    # Check that local_device_ids was passed correctly as a list of integers
+    _, kwargs = mock_init.call_args
+    self.assertEqual(kwargs["local_device_ids"], [0, 2, 3])
+    self.assertEqual(kwargs["coordinator_address"], "10.0.0.1:1234")
+
+  @mock.patch.dict(
+      os.environ,
+      {
+          "JAX_COORDINATOR_IP": "10.0.0.1",
+          "JAX_COORDINATOR_PORT": "1234",
+          "NNODES": "1",
+          "NODE_RANK": "0",
+          "CUDA_VISIBLE_DEVICES": "GPU-8f2e3072-...",  # Invalid format for integer parsing
+      },
+  )
+  @mock.patch("jax.distributed.initialize")
+  @mock.patch("jax.devices")
+  @mock.patch("maxtext.utils.max_logging.log")
+  def test_initialize_jax_for_gpu_invalid_devices(self, _mock_log, mock_devices, mock_init):
+    """Verifies fallback behavior when parsing fails (e.g., UUIDs)."""
+    raw_keys = {"jax_distributed_initialization_timeout": 300}
+    max_utils.initialize_jax_for_gpu(raw_keys)
+    # Check that it falls back to None (JAX auto-detection default) on error
+    _, kwargs = mock_init.call_args
+    self.assertIsNone(kwargs.get("local_device_ids"))
+    self.assertEqual(kwargs["coordinator_address"], "10.0.0.1:1234")
+
+  @mock.patch.dict(
+      os.environ,
+      {
+          "JAX_COORDINATOR_IP": "10.0.0.1",
+          "JAX_COORDINATOR_PORT": "1234",
+          "NNODES": "1",
+          "NODE_RANK": "0",
+      },
+  )
+  @mock.patch("jax.distributed.initialize")
+  @mock.patch("jax.devices")
+  @mock.patch("maxtext.utils.max_logging.log")
+  def test_initialize_jax_for_gpu_no_devices(self, _mock_log, mock_devices, mock_init):
+    """Verifies that no error occurs when CUDA_VISIBLE_DEVICES is not set"""
+    raw_keys = {"jax_distributed_initialization_timeout": 300}
+    max_utils.initialize_jax_for_gpu(raw_keys)
+    _, kwargs = mock_init.call_args
+    self.assertIsNone(kwargs.get("local_device_ids"))
+    self.assertEqual(kwargs["coordinator_address"], "10.0.0.1:1234")
+
+
 if __name__ == "__main__":
   unittest.main()