Repository: google/jax
Pull request #22592 (2024-07-23 12:14:45)

mesh_utils: allow meshes that do not include device at (0, 0, 0).
mesh_utils: allow meshes that do not include device at (0, 0, 0).

This is required to allow the use of subslices: e.g., the two halves
of a TPU slice.  One of them will not include the device at
coordinates (0, 0, 0).

E.g., assume we have a TPU v4 1x2x1 slice.

BEFORE THIS CL, if we call _get_physical_tpu_mesh() (an auxiliary for
the public create_device_mesh()) with

jax_devices=[device(0,TPU_DEVICE,coords=[0,0,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3)]

we get the expected result

[[[device(0,TPU_DEVICE,coords=[0,0,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3)]]]

However, if we call it with

jax_devices=[device(1,TPU_DEVICE,coords=[0,1,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3)]

we get the wrong mesh

[[[None]
  [device(1,TPU_DEVICE,coords=[0,1,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3)]]]

That's because the code before this CL assumed the the incoming
jax_devices are arranged in a cuboid that starts at (0, 0, 0).  When
working with subslices (e.g., half of a TPU slice) that is not always
the case.

AFTER THIS CL, the second case will return
[[[device(1,TPU_DEVICE,coords=[0,1,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3)]]]

For each dimension from the TPU coordinates, this CL computes the min
/ max; we expect the provided devices to fill the [min, max] interval
(in that dimension).  By requesting this for each dimension, we
request that the set of provided devices constitute a cuboid, but,
unlike before this CL, that cuboid does not need to include (0, 0, 0):
it can be "translated", which allows e.g., both half-slices of a big
slice.



Patch:
diff --git a/jax/experimental/mesh_utils.py b/jax/experimental/mesh_utils.py
index dd112db3b..473700024 100644
--- a/jax/experimental/mesh_utils.py
+++ b/jax/experimental/mesh_utils.py
@@ -575,7 +575,12 @@ def _get_physical_tpu_mesh(jax_devices: Sequence[Any]) -> np.ndarray:
 
   Args:
     jax_devices: A list of JAX devices in a TPU slice in process-tiled z, y, x,
-      core order, e.g. from jax.devices().
+      core order, e.g. from jax.devices(). The coordinates of these devices
+      should constitute a cuboid with no holes; e.g., the coordinates can be
+      {(1, 0, 0), (1, 0, 1), (1, 1, 0), (1, 1, 1)} (a 1x2x2 cuboid); passing
+      only 3 of these devices would result in a "hole" in that cuboid, which is
+      an error.  As in our example, the cuboid is not required to include the
+      point (0, 0, 0).
 
   Returns:
     A np.ndarray of JAX devices with shape [global_x, global_y, global_z]. On
@@ -583,24 +588,57 @@ def _get_physical_tpu_mesh(jax_devices: Sequence[Any]) -> np.ndarray:
   """
   device_kind = jax_devices[0].device_kind
   device_coords = [d.coords for d in jax_devices]
-  dims = tuple(d + 1 for d in max(device_coords))
+  coord_size = len(device_coords[0])
+  # Position-wise max and min coordinates:
+  max_coords = tuple(
+      max(dc[i] for dc in device_coords) for i in range(coord_size)
+  )
+  min_coords = tuple(
+      min(dc[i] for dc in device_coords) for i in range(coord_size)
+  )
+  dims = tuple(h - l + 1 for (h, l) in zip(max_coords, min_coords))
+
+  max_cores_per_chip = max(d.core_on_chip for d in jax_devices)
+  min_cores_per_chip = min(d.core_on_chip for d in jax_devices)
+  cores_per_chip = max_cores_per_chip - min_cores_per_chip + 1
+
   assert len(dims) == 3, dims
+  assert (
+      len(jax_devices) == np.prod(dims) * cores_per_chip
+  ), f'{jax_devices=} {dims=} {cores_per_chip=}'
+
   if device_kind in (_TPU_V2, _TPU_V3):
-    cores_per_chip = max(d.core_on_chip for d in jax_devices) + 1
     out = np.empty(dims[:2] + (cores_per_chip,), dtype=object)
-    for coords, d in zip(device_coords, jax_devices):
+    for d in jax_devices:
+      coords = d.coords
       assert coords[2] == 0, d
-      out[coords[0], coords[1], d.core_on_chip] = d
+      out[
+          coords[0] - min_coords[0],
+          coords[1] - min_coords[1],
+          d.core_on_chip - min_cores_per_chip,
+      ] = d
   else:
     out = np.empty(dims, dtype=object)
-    for coords, d in zip(device_coords, jax_devices):
+    for d in jax_devices:
+      coords = d.coords
       if d.core_on_chip != 0:
         raise AssertionError(
             'Creating meshes for TPU >v3 requires one device per chip'
             f' ("megacore" mode). Got device id {d.core_on_chip} for a device'
             f' of kind {device_kind}: {d}.'
         )
-      out[coords[0], coords[1], coords[2]] = d
+      out[
+          coords[0] - min_coords[0],
+          coords[1] - min_coords[1],
+          coords[2] - min_coords[2],
+      ] = d
+
+  # Check there is no "hole" in the mesh we constructed.
+  if (out == None).any():  # pylint: disable=singleton-comparison
+    raise AssertionError(
+        'Constructed mesh contains a "hole"; probable cause: coordinates '
+        f'of jax_devices are not a contiguous cuboid: {jax_devices}'
+    )
   return out
 
 


Test patch:
diff --git a/tests/mesh_utils_test.py b/tests/mesh_utils_test.py
index 1c8893f36..6916c2c37 100644
--- a/tests/mesh_utils_test.py
+++ b/tests/mesh_utils_test.py
@@ -54,9 +54,18 @@ class MockTpuDevice:
 def mock_tpu_devices(x, y, z, dev_kind, one_device_per_chip, num_slices=1,
                      reorder=False):
   """Produce fake jax.devices() output for a TPU slice."""
+  assert x > 0 and y > 0 and z > 0
+
   cores_per_chip = 1 if one_device_per_chip else 2
-  nxd, nyd, nzd = (2, 2, 1)
+
+  # 3D shape of the mesh of devices on each host (= process).
+  nxd, nyd, nzd = (min(x, 2), min(y, 2), 1)
+  # 3D shape of the mesh of hosts (= processes):
   nxp, nyp, nzp = x // nxd, y // nyd, z // nzd
+  assert nxp * nxd == x
+  assert nyp * nyd == y
+  assert nzp * nzd == z
+
   def mock_tpu_device(core_on_chip, xd, yd, zd, xp, yp, zp, slice_index):
     process_index = xp + nxp * (yp + nyp * (zp + nzp * slice_index))
     coords =  (xd + nxd * xp, yd + nyd * yp, zd + nzd * zp)
@@ -71,23 +80,46 @@ def mock_tpu_devices(x, y, z, dev_kind, one_device_per_chip, num_slices=1,
              for core_on_chip in range(cores_per_chip)]
   if reorder:
     devices = devices[::-1]
-  _validate_mocked_process_indices(devices, one_device_per_chip)
+
+  # Validate the generated mock devices:
+  num_local_chips = nxd * nyd  # Number of mock devices / process.
+  if num_local_chips < 4:
+    # Sub-host slice = fewer than the 4 chips available on a host:
+    # e.g., 1x1 TPU v2.  All devices should be on one host.
+    num_all_chips = x * y * z
+    assert num_all_chips == num_local_chips, f'Bad shape: {x=}, {y=}, {z=}'
+    # Implied by the previous assertion, but let's be explicit:
+    assert z == 1
+    _validate_mocked_devices_for_subhost_slice(devices, x, y, cores_per_chip)
+  else:
+    _validate_mocked_devices(devices, num_local_chips * cores_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):
+def _validate_mocked_devices_for_subhost_slice(devices, x, y, cores_per_chip):
+  first_device = devices[0]
+  distinct_coords = set()
+  for d in devices:
+    assert d.process_index == first_device.process_index
+    assert d.coords[0] >= 0 and d.coords[0] < x
+    assert d.coords[1] >= 0 and d.coords[1] < y
+    assert d.coords[2] == 0
+    assert d.core_on_chip >= 0 and d.core_on_chip < cores_per_chip
+    distinct_coords.add((d.coords[0], d.coords[1], 0, d.core_on_chip))
+  assert len(distinct_coords) == x * y * cores_per_chip
+
+
+# If this function raises, it's a bug in the test code!
+def _validate_mocked_devices(devices, num_local_devices):
+  # NOTE: this function is not called for sub-host slices.
   process_to_devices = collections.defaultdict(list)
   for d in devices:
     process_to_devices[d.process_index].append(d)
 
   for local_devices in process_to_devices.values():
-    if one_device_per_chip:
-      # 4 devices per process
-      assert len(local_devices) == 4, local_devices
-    else:
-      # 8 devices per process
-      assert len(local_devices) == 8, local_devices
+    assert len(local_devices) == num_local_devices, local_devices
     # All devices have same z coord
     assert len({d.coords[2] for d in local_devices}) == 1, local_devices
     # All devices in a 2x2 subgrid
@@ -98,6 +130,11 @@ def _validate_mocked_process_indices(devices, one_device_per_chip):
     assert {d.coords for d in local_devices} == expected, local_devices
 
 
+def mock_1x1_devices():
+  """Hard-coded reproduction of jax.devices() output on v3-1x1."""
+  return mock_tpu_devices(1, 1, 1, 'TPU v3', False)
+
+
 def mock_2x2_devices():
   """Hard-coded reproduction of jax.devices() output on v3-2x2."""
   return mock_tpu_devices(2, 2, 1, 'TPU v3', False)
@@ -113,6 +150,11 @@ def mock_8x8_devices(one_device_per_chip=False):
   return mock_tpu_devices(8, 8, 1, 'TPU v3', one_device_per_chip)
 
 
+def mock_1x2x1_devices(one_device_per_chip):
+  """Hard-coded reproduction of jax.devices() output on 2x2x1."""
+  return mock_tpu_devices(1, 2, 1, 'TPU v4', one_device_per_chip)
+
+
 def mock_2x2x1_devices(one_device_per_chip):
   """Hard-coded reproduction of jax.devices() output on 2x2x1."""
   return mock_tpu_devices(2, 2, 1, 'TPU v4', one_device_per_chip)
@@ -156,18 +198,32 @@ def mock_8x8x16_devices(one_device_per_chip):
 class MeshUtilsTest(test_util.JaxTestCase):
 
   @parameterized.named_parameters(
+      ('1x1', mock_1x1_devices, (1, 1, 1, 2)),
+      ('2x2', mock_2x2_devices, (2, 2, 1, 2)),
+      ('4x4', mock_4x4_devices, (4, 4, 1, 2)),
+      ('8x8', mock_8x8_devices, (8, 8, 1, 2)),
+  )
+  def test_bounds_from_last_device_2d(self, devices, expected_bounds):
+    self.assertEqual(
+        mesh_utils._bounds_from_last_device(devices()[-1]),
+        expected_bounds)
+
+  @parameterized.named_parameters(
+      ('1x2x1_t', mock_1x2x1_devices, True, (1, 2, 1, 1)),
+      ('1x2x1_f', mock_1x2x1_devices, False, (1, 2, 1, 2)),
       ('2x2x1_t', mock_2x2x1_devices, True, (2, 2, 1, 1)),
       ('2x2x1_f', mock_2x2x1_devices, False, (2, 2, 1, 2)),
       ('8x8x16_t', mock_8x8x16_devices, True, (8, 8, 16, 1)),
       ('8x8x16_f', mock_8x8x16_devices, False, (8, 8, 16, 2)),
   )
-  def test_bounds_from_last_device(self, devices, one_device_per_chip,
-                                   expected_bounds):
+  def test_bounds_from_last_device_3d(self, devices, one_device_per_chip,
+                                      expected_bounds):
     self.assertEqual(
         mesh_utils._bounds_from_last_device(devices(one_device_per_chip)[-1]),
         expected_bounds)
 
   @parameterized.named_parameters(
+      ('1x2x1_t', (1, 2, 1), True),
       ('4x4x4_t', (4, 4, 4), True),
       ('4x4x4_f', (4, 4, 4), False),
       ('8x8x16_t', (8, 8, 16), True),
@@ -184,6 +240,102 @@ class MeshUtilsTest(test_util.JaxTestCase):
         for k in range(z):
           self.assertEqual(normalized[i, j, k].coords, (i, j, k))
 
+  def test_get_physical_tpu_mesh_with_subslice_TPU_v2_1x1(self):
+    one_device_per_chip = False  # Each TPU v2 chip has 2 devices.
+    device_list = mock_tpu_devices(1, 1, 1, 'TPU v2', one_device_per_chip)
+    device_array = mesh_utils._get_physical_tpu_mesh(device_list)
+    self.assertEqual(device_array.shape, (1, 1, 2))
+
+    # A subslice that includes the device at (0, 0, 0): core #0 of the
+    # device at (x, y, z) == (0, 0, 0).
+    subslice0 = mesh_utils._get_physical_tpu_mesh([device_array[0, 0, 0]])
+    self.assertEqual(subslice0.shape, (1, 1, 1))
+    self.assertEqual(subslice0[0, 0, 0], device_array[0, 0, 0])
+    self.assertEqual(subslice0[0, 0, 0].coords, (0, 0, 0))
+    self.assertEqual(subslice0[0, 0, 0].core_on_chip, 0)
+
+    # Another subsublice, without the device at (0, 0, 0): core #1 of
+    # the device at (x, y, z) == (0, 0, 0).
+    subslice1 = mesh_utils._get_physical_tpu_mesh([device_array[0, 0, 1]])
+    self.assertEqual(subslice1.shape, (1, 1, 1))
+    self.assertEqual(subslice1[0, 0, 0], device_array[0, 0, 1])
+    self.assertEqual(subslice1[0, 0, 0].coords, (0, 0, 0))
+    self.assertEqual(subslice1[0, 0, 0].core_on_chip, 1)
+
+  def test_get_physical_tpu_mesh_with_subslice_TPU_v4_1x2x1(self):
+    one_device_per_chip = True  # For TPU v4, chip == device.
+    device_list = mock_tpu_devices(1, 2, 1, 'TPU v4', one_device_per_chip)
+    device_array = mesh_utils._get_physical_tpu_mesh(device_list)
+    self.assertEqual(device_array.shape, (1, 2, 1))
+
+    # A subslice that includes the device at (0, 0, 0).
+    subslice0 = mesh_utils._get_physical_tpu_mesh([device_array[0, 0, 0]])
+    self.assertEqual(subslice0.shape, (1, 1, 1))
+    self.assertEqual(subslice0[0, 0, 0], device_array[0, 0, 0])
+
+    # Another subsublice, without the device at (0, 0, 0).
+    subslice1 = mesh_utils._get_physical_tpu_mesh([device_array[0, 1, 0]])
+    self.assertEqual(subslice1.shape, (1, 1, 1))
+    self.assertEqual(subslice1[0, 0, 0], device_array[0, 1, 0])
+
+  def test_get_physical_tpu_mesh_with_subslice_TPU_v5e_4x4(self):
+    one_device_per_chip = True  # For TPU v5e, chip == device.
+    device_list = mock_tpu_devices(4, 4, 1, 'TPU v5e', one_device_per_chip)
+    device_array = mesh_utils._get_physical_tpu_mesh(device_list)
+
+    # `device_array` is isomorphic with a 4x4 grid (z coord == 0).
+    self.assertEqual(device_array.shape, (4, 4, 1))
+
+    # Two subslices: each subslice has shape (4, 2); first one starts
+    # at (x=0, y=0), the other at (x=0, y=2); visually, the left
+    # and right halves of the (4, 4) grid.
+    for start_y in (0, 2):
+      subslice_devices = []
+      for x in range(4):
+        for delta_y in range(2):
+          subslice_devices.append(device_array[x, start_y + delta_y, 0])
+      logging.info(
+          'start_y=%s subslice_devices=%s', start_y, subslice_devices
+      )
+      subslice = mesh_utils._get_physical_tpu_mesh(subslice_devices)
+      self.assertEqual(subslice.shape, (4, 2, 1))
+      for x in range(4):
+        for delta_y in range(2):
+          self.assertEqual(
+              subslice[x, delta_y],
+              device_array[x, start_y + delta_y, 0],
+          )
+
+  def test_get_physical_tpu_mesh_with_bad_subslice(self):
+    one_device_per_chip = True  # For TPU v5e, chip == device.
+    device_list = mock_tpu_devices(4, 4, 1, 'TPU v5e', one_device_per_chip)
+    device_array = mesh_utils._get_physical_tpu_mesh(device_list)
+
+    self.assertEqual(device_array.shape, (4, 4, 1))
+
+    # Second subslice from
+    # test_get_physical_tpu_mesh_with_subslice_TPU_v5e_4x4, without
+    # the device from its top-left corner (device from (0, 2, 0)).
+    start_y = 2
+    subslice_devices = []
+    for x in range(4):
+      for delta_y in range(2):
+        if (x == 0) and (delta_y == 0):
+          # Skip device from (0, 2, 0).
+          continue
+        subslice_devices.append(device_array[x, start_y + delta_y, 0])
+
+    # subslice_devices are obviously not a cuboid: only 7 devices.
+    with self.assertRaises(AssertionError):
+      mesh_utils._get_physical_tpu_mesh(subslice_devices)
+
+    # Make it a bit harder, such that just a simple test on
+    # len(subslice_devices) is not enough: 8 devices, but two of them
+    # are identical (device from (2, 2, 0) is duplicated.
+    subslice_devices.append(device_array[2, 2, 0])
+    with self.assertRaisesRegex(AssertionError, 'not a contiguous cuboid'):
+      mesh_utils._get_physical_tpu_mesh(subslice_devices)
+
   @parameterized.named_parameters(
       ('2x2x1', mock_2x2x1_devices, [1, 1, 4], [(), (), (0, 1, 2)]),
       ('2x2x4', mock_2x2x4_devices, [1, 4, 4], [(), (2,), (0, 1)]),