Repository: google/jax
Pull request #18940 (2023-12-12 16:41:53)

`mesh_utils.create_hybrid_device_mesh`: make sorting granules by key user configurable.
`mesh_utils.create_hybrid_device_mesh`: make sorting granules by key user configurable.

When sorting by granule key is disabled, the granules are used to create the mesh in the order in which they appear in the sequence of devices.


google-cla[bot]: Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this [failed invocation](https://github.com/google/jax/pull/18940/checks?check_run_id=19565867858) of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

Patch:
diff --git a/jax/experimental/mesh_utils.py b/jax/experimental/mesh_utils.py
index faf48a841..8b05e511e 100644
--- a/jax/experimental/mesh_utils.py
+++ b/jax/experimental/mesh_utils.py
@@ -321,10 +321,13 @@ def create_device_mesh(
     device_mesh = np.asarray(devices).reshape(mesh_shape)
     return device_mesh
 
-def create_hybrid_device_mesh(mesh_shape: Sequence[int],
-                              dcn_mesh_shape: Sequence[int],
-                              devices: Optional[Sequence[Any]] = None, *,
-                              process_is_granule: bool = False) -> np.ndarray:
+def create_hybrid_device_mesh(
+    mesh_shape: Sequence[int],
+    dcn_mesh_shape: Sequence[int],
+    devices: Optional[Sequence[Any]] = None, *,
+    process_is_granule: bool = False,
+    should_sort_granules_by_key: bool = True,
+) -> np.ndarray:
   """Creates a device mesh for hybrid (e.g., ICI and DCN) parallelism.
 
   Args:
@@ -339,6 +342,9 @@ def create_hybrid_device_mesh(mesh_shape: Sequence[int],
       of the slower/outer network. Otherwise it will look for slice_index
       attributes on devices and use slices as the units. Enabling this is meant
       as a fallback for platforms (e.g., GPU) that don't set slice_index.
+    should_sort_granules_by_key: Whether device granules should be sorted by the
+      granule key, either slice or process index, depending on
+      process_is_granule.
 
   Raises:
     ValueError: if the number of slices to which the `devices` belong doesn't
@@ -356,7 +362,10 @@ def create_hybrid_device_mesh(mesh_shape: Sequence[int],
   granule_dict = collections.defaultdict(list)
   for dev in devices:
     granule_dict[getattr(dev, attr)].append(dev)
-  granules = [granule_dict[key] for key in sorted(granule_dict.keys())]
+  granules = (
+      [granule_dict[key] for key in sorted(granule_dict.keys())]
+      if should_sort_granules_by_key
+      else granule_dict.values())
   if np.prod(dcn_mesh_shape) != len(granules):
     raise ValueError(
         f'Number of slices {len(granules)} must equal the product of '


Test patch:
diff --git a/tests/mesh_utils_test.py b/tests/mesh_utils_test.py
index a4b1fda5c..0ed5ce4c5 100644
--- a/tests/mesh_utils_test.py
+++ b/tests/mesh_utils_test.py
@@ -63,6 +63,7 @@ def mock_tpu_devices(x, y, z, dev_kind, one_device_per_chip, num_slices=1,
   _validate_mocked_process_indices(devices, one_device_per_chip)
   return devices
 
+
 # If this function raises, it's a bug in the test code!
 def _validate_mocked_process_indices(devices, one_device_per_chip):
   process_to_devices = collections.defaultdict(list)
@@ -202,7 +203,48 @@ class MeshUtilsTest(test_util.JaxTestCase):
         mesh_shape, dcn_mesh_shape, devices)
     total_mesh_shape = tuple(
         m1 * m2 for m1, m2 in zip(mesh_shape, dcn_mesh_shape))
-    assert mesh.shape == total_mesh_shape
+    self.assertEqual(mesh.shape, total_mesh_shape)
+
+  @parameterized.named_parameters(
+      ('2X4x4x4a', (1, 16, 4), (2, 1, 1)),
+      ('2X4x4x4b', (1, 4, 16), (1, 2, 1)),
+  )
+  def test_create_hybrid_device_mesh_device_sorting(
+      self,
+      mesh_shape: tuple[int, ...],
+      dcn_mesh_shape: tuple[int, ...],
+  ):
+    devices = mock_tpu_devices(4, 4, 4, 'TPU v4', True, 2)
+    reversed_slices_devices = list(
+        np.flip(np.array(devices).reshape(2, -1), axis=0).flat)
+    mesh = mesh_utils.create_hybrid_device_mesh(
+        mesh_shape,
+        dcn_mesh_shape,
+        devices,
+        should_sort_granules_by_key=False,
+    )
+    sorted_slices_mesh = mesh_utils.create_hybrid_device_mesh(
+        mesh_shape,
+        dcn_mesh_shape,
+        reversed_slices_devices,
+        should_sort_granules_by_key=True,
+    )
+    np.testing.assert_array_equal(mesh, sorted_slices_mesh)
+    self.assertSetEqual(
+        {0, 1},
+        {d.slice_index for d in sorted_slices_mesh.flat},
+    )
+
+    reversed_slices_mesh = mesh_utils.create_hybrid_device_mesh(
+        mesh_shape,
+        dcn_mesh_shape,
+        reversed_slices_devices,
+        should_sort_granules_by_key=False,
+    )
+    self.assertSetEqual(
+        {1, 0},
+        {d.slice_index for d in reversed_slices_mesh.flat},
+    )
 
   @parameterized.named_parameters(
       # Physical ring order over tray