Repository: google/jax
Pull request #21705 (2024-06-06 16:15:13)

Reverts dfe61285093ff826e1ad23bb36b77a42c01040b4
Reverts dfe61285093ff826e1ad23bb36b77a42c01040b4



Patch:
diff --git a/jax/_src/array.py b/jax/_src/array.py
index 854975c97..d2064a16e 100644
--- a/jax/_src/array.py
+++ b/jax/_src/array.py
@@ -43,7 +43,7 @@ from jax._src.lib import xla_extension as xe
 from jax._src.sharding import Sharding
 from jax._src.sharding_impls import (
     PmapSharding, SingleDeviceSharding,
-    device_replica_id_map, hashed_index, num_addressable_indices)  # pyformat: disable
+    device_replica_id_map, hashed_index, num_addressable_indices, local_to_global_shape)  # pyformat: disable
 from jax._src.typing import ArrayLike, DLDeviceType
 from jax._src.util import safe_zip, unzip3, use_cpp_class, use_cpp_method
 import numpy as np
@@ -734,7 +734,7 @@ def make_array_from_callback(
 def make_array_from_process_local_data(
     sharding: Sharding,
     local_data: np.ndarray,
-    global_shape: tuple[int, ...],
+    global_shape: Shape | None = None,
 ) -> ArrayImpl:
   # pyformat: disable
   """Creates distributed tensor using the data available in process.
@@ -743,26 +743,38 @@ def make_array_from_process_local_data(
   assumes that the data is available in the process and takes care of the
   index wrangling.
 
-  Note, if the two hosts are replicas, host_local_data should be identical as
-  well.
-  Each dimension of the shape of host_local_data should either match
-  global_shape or the # indices the devices on this process need to
-  address. For example if dimension $i$ is fully sharded then this size would be
-  `per_device_shape[i] * jax.local_device_count()`.
+  The most common case is when the sharding is sharded across the batch
+  dimension and each host just loads its corresponding sub-batch. This function
+  supports more general cases as well, such as mixed multi-host and multi-axis
+  replication and sharding but you would need to compute the size and the
+  contents of process-local data correctly to satisfy the sharding constraints.
 
-  If the shape matches global shape, each device slice will just lookup
-  the slice in the local_data. In the latter case the global slice of each
-  device will be mapped into local slice of `local_data` array. For example,
-  if given process only addresses slices (8, 12) and  (24, 28), then
-  these slices will be mapped into (0, 4) and (4, 8) of the `local_data`.
+  In particular, if any two hosts are replicas, host_local_data should be
+  identical as well.
 
-  This function can be used to create tensors from dataset feeding pipelines.
+  The global_shape is optional. If not provided it will be be inferred from
+  the local_data and sharding, under the assumption that
+  each host represents only their own data for uniform sharding. If sharding
+  is non-uniform, (see note below) an exception will be raised.
 
-  The most common case is when the sharding is fully sharded across the batch
-  dimension and each host just loads its corresponding sub-batch. This function
-  supports more general case as well, such as multi-host replication
-  but you would need to compute the size and the contents of process-local data
-  correctly to satisfy the replication constraints.
+  Setting global_shape explicitly allows for finer grain control and works with
+  non-uniform shardings. Each dimension of global_shape must either match
+  host_local_data, or match the inferred global shape of the sharding (in which
+  case it is equivalent to setting it to None, but is more explicit).
+
+  For example if dimension `i` is fully sharded then this size would be
+  `per_device_shape[i] * jax.local_device_count()`.  Each device will be mapped
+  into local slice of `local_data` array. For example, if given process
+  addresses slices (8, 12) and  (24, 28), then these slices will be mapped
+  into (0, 4) and (4, 8) of the `local_data`.
+
+  For each dimension where global_shapes matches local_shape, each device
+  will lookup the slice in the local_data. For example if
+  global_shape == local_data.shape, the local data is assumed to be the
+  actual target array that will be sharded into device.
+
+  If global_shape is the same as local_data.shape, then the data must
+  be the same across all hosts.
 
   Examples:
     >>> from jax.sharding import PartitionSpec as P
@@ -784,19 +796,71 @@ def make_array_from_process_local_data(
     >>> assert output_global_array.addressable_data(0).shape == per_device_shape
     >>> assert output_global_array.shape == global_shape
 
+  NB: While most shardings are uniform, It is possible to design am exotic
+  sharding mesh where each process's  devices will be arranged in a non-grid
+  like pattern in some dimensions, or for indices to overlap non-trivially.
+  Such sharding is called "non-uniform" in those dimensions. In that case,
+  the global shape along those directions must match local shape as there is
+  no meaningful way to represent all needed
+  per-process data in non-overlapping fashion. For example for global_shape 4x4
+  if sharding looks like this:
+
+      0123
+      2103
+      4675
+      4567
+
+  with 4 processes, containing devices (0,1), (2, 3), (4, 5), (6, 7) respectively.
+  Then the data for each host look like
+
+      xx..    ..xx     ....    ....
+      .xx.    x..x     ....    ....
+      ....    ....     x..x    .xx.
+      ....    ....     xx..    ..xx
+
+  the sharding is uniform on rows (each host requires either rows 1-2, or rows 3-4)
+  and non-uniform on columns (hosts require overlapping but not matching
+  set of columns). Thus local data must have the shape 2x4 or 4x4
+  for all hosts, even though each  host can potentially fit into 2x2 shape.
+  In this case user must provide global_shape explicitly and for
+  local_shape=(2, 4), potentially valid global shapes are (2, 4) and (4, 4).
+
+  On the other hand for sharding:
+
+      0213   x.x.  .x.x.  ....  ....
+      0213   x.x.  .x.x.  ....  ....
+      4657   ....  ....   .x.x  x.x.
+      4657   ....  ....   .x.x  x.x.
+
+  for local_shape=(2, 2) this function can accept a choice of 2x2, 2x4, 4x2
+  and 4x4 global shapes. Setting global_shape to None, is equivalent to
+  setting it to (4, 4) in this case.
+
   Args:
     sharding: sharding of the global tensor.
     host_local_data: data on the host to be placed on local devices. Each
       dimension should either match global_shape, or match
       num_addressable_indices(dim).
-    global_shape: the target shape of the global tensor. In some cases this
-      parameter can be inferred from sharding and host_local_data, however it is
-      useful to catch common sharding errors.
+    global_shape: the target shape of the global tensor. If None,
+      will infer from host_local_data and sharding.
 
   Returns:
-    Tensor that will have sharding=sharding.
+    Tensor that will have sharding=sharding and of shape global_shape.
   """
   # pyformat: enable
+  # TODO(sandler): consider supporting partially specified global_shape or
+  # making local_to_global_shape available in the api.
+  local_shape = local_data.shape
+  if global_shape is None:
+    global_shape = local_to_global_shape(sharding, local_shape)  # type: ignore[assignment]
+    assert global_shape is not None
+    if None in global_shape:
+      raise ValueError(
+          "Unable to compute global_shape due to non-uniform sharding."
+          f" Specify global shape directly. Partially computed {global_shape=}."
+      )
+  elif None in global_shape:
+    raise ValueError(f"{global_shape=} has Nones. This is not supported.")
   shard_shape = sharding.shard_shape(global_shape)
   full_dim = []
   for i, (data_dim, global_dim) in enumerate(
diff --git a/jax/_src/sharding_impls.py b/jax/_src/sharding_impls.py
index 3824eb21e..3e98327b4 100644
--- a/jax/_src/sharding_impls.py
+++ b/jax/_src/sharding_impls.py
@@ -1301,6 +1301,192 @@ def _slice_as_tuple(s: slice):
   return (s.start, s.stop)
 
 
+class NonUniformShardingError(ValueError):
+  """Raised when sharding is not uniform across processes."""
+
+
+def get_process_index_and_count(
+    tensor_sharding: sharding.Sharding,
+    dim: int,
+    ndims: int,
+) -> tuple[int, int]:
+  """Get current process index and number of unique processes for given dimension.
+
+  This function facilitates mapping of process-level data to individual
+  devices. Each process can use its index to obtain the data corresponding
+  to that index. If process level data is sharded on multiple dimensions
+  this function can be used to build the cross product of indices in
+  each sharded axis. Processes that need to load the same data will have
+  the same index. For shardings whose per-process data is not distributed
+  on a grid, the number of distinct shards will be such that it is possible to
+  build the target shape while maintaining a "cube" shape of local-process data.
+
+  For example, in case of 4 hosts with sharding distributed like so:
+
+  1234
+  2143
+
+  For dim 0 (rows): all processes need to access all rows, so we return (0, 1)
+  For dim 1 (cols):
+     process 1 and 2 returns index 0 out of 2 (need cols 0 and 1),
+     process 3 and 4 returns index 1 out of 2 (need cols 2 and 3).
+
+  On the other hand, for a sharding like:
+
+  1212
+  3434
+
+  Dim 0 (rows): process 1 and 2 returns (0, 2), process 3 and 4 returns (1, 2)
+  Dim 1 (cols): process 1 and 3 returns (0, 2), process 2 and 4 returns (1, 2)
+
+  Note: This function requires sharding to be process uniform in dimension
+  `dim`:
+   each process has the same number of addressable indices in that
+  dimension and all index sets across processes are either disjoint or the same.
+
+  For sharding to be process uniform the addressable shards doesn't need to
+  form contiguous subtensor, or even a sparse grid  and  in case of
+  interleaved high-dimensional tensor it is possible for sharding to be
+  process uniform only in some dimensions but not others.
+
+  For example:
+    1111 and 12 and 1212 and 1212
+    2222     21     2121     1212
+
+  are all sharding uniform, in both dimensions. However
+
+    1122
+    2121
+    1121
+    1222
+
+  is uniform in dimension 0 (both hosts access all rows), but
+  is not uniform in dimension 1 (host 1 accesses columns: 0, 1, and 3),
+  while host 2 accesses (0, 1, 2, 3).
+
+  Returns:
+    A tuple of (index, num_distinct_shards) for the given dimension.
+    It is guaranteed that `index` will cover 0 to `num_distinct_shards - 1`,
+    across all processes.
+
+  Raises:
+    NonUniformShardingError: if the sharding is not process uniform in dimension
+    `dim`.
+  """
+  # TODO(sandler, yashkatariya): Consider making this function public.
+
+  if (
+      tensor_sharding.is_fully_addressable
+      or tensor_sharding.is_fully_replicated
+  ):
+    return (0, 1)
+  num_devices = len(tensor_sharding.device_set)
+  # Get device to indices map, we don't care about the concrete
+  # global shape here, only to get the distribution of shards across the tensor
+  # using (num_devices, num_devices, ...)  This is a universal shape that is
+  # compatible with any mesh with num_devices.
+  device_map = tensor_sharding.devices_indices_map((num_devices,) * ndims)
+
+  # Get the slices for 'dim' for all devices.
+  global_slice = {k: v[dim] for k, v in device_map.items()}
+
+  # Contains mapping from process_index to a set of slices for that process.
+  process_to_slice = collections.defaultdict(set)
+  # Contains global set of slices across all processes.
+  all_slices = set()
+
+  # Compute the set of slices for each process and the global set of slices.
+  for d, v in global_slice.items():
+    key = (v.start, v.stop)
+    process_to_slice[d.process_index].add(key)
+    all_slices.add(key)
+
+  # Get the set of slices for the current process which we will use to compute
+  # the index of the current process.
+  current_pid = next(iter(tensor_sharding.addressable_devices)).process_index
+  addressable_slices = frozenset(process_to_slice[current_pid])
+
+  # Verify that all processes have the same number of slices.
+  slices_per_process = len(addressable_slices)
+  if any(len(x) != slices_per_process for x in process_to_slice.values()):
+    raise NonUniformShardingError(
+        f'{tensor_sharding=} is non-uniform on {dim=} as some processes have '
+        'different number of slices.'
+    )
+  unique_processes = list({frozenset(x) for x in process_to_slice.values()})
+
+  # After removing duplicate processes all unique slices should
+  # cover the dimension exactly once. If they don' it means that
+  # the sharding is not uniform.
+  if sum(len(h) for h in unique_processes) != len(all_slices):
+    raise NonUniformShardingError(
+        f'{tensor_sharding=} is non-uniform on {dim=}'
+    )
+  return (unique_processes.index(addressable_slices), len(unique_processes))
+
+
+def local_to_global_shape(
+    sharding: sharding.Sharding,
+    local_shape: Shape,
+) -> tuple[int | None, ...]:
+  """Computes the global shape given the per process if possible.
+
+  The returned shape will have the size of the global tensor in that dimension
+  or None, if it is not computable. The latter can happen when sharding
+  is not uniform along that dimension, e.g. different hosts require
+  different shapes, or if different processes have partial data overlap.
+
+  If at most one dimension is sharded the shape is always computable.
+  Generally, global shape is computable for most practical meshes (including
+  topology aware such as meshes returned by mesh_utils.create_device_mesh)
+
+  Some examples: Suppose mesh is {'a': 2, 'b': 2, 'c': 2} with 2 devices
+  per host, 4 hosts total. For different specs we get:
+  - P():
+      global_shape = local_shape
+
+  - P(('a', 'b', 'c'), None):
+      global_shape =  (4 * local_shape[0], local_shape[1])
+      Note: per device shape is (local_shape[0] / 2, local_shape[1])
+
+  - P(('a', 'b'), None)
+      global_shape =  (4 * local_shape[0], local_shape[1])
+      # NB: the same global shape as above, since sharding along 'c' dimension
+      # happens to be within process, and thus doesn't affect the global shape.
+      # The underlying difference will be in the per *device* shape, which
+      # would be  (local_shape[0], local_shape[1]) in this case.
+
+  - P(None, ('a', 'c'))
+      global_shape = (local_shape[0], 2 * local_shape[1])
+      # Per device shape is (local_shape[0], local_shape[1] / 2)
+  - P(('a', 'c'), 'b'):
+      global_shape = (2 * local_shape[0], 2 * local_shape[1])
+      # Per device shape is (local_shape[0] / 2, local_shape[1])
+  - If devices in the Mesh are randomly permuted: For any partition spec
+  which shards more than 1 axis:  e.g. P('a', ('b', 'c')):
+      global_shape = (None, None)
+
+  Args:
+    local_shape: global shape of the tensor.
+
+  Returns:
+    global_shape with Nones in non-uniform dimensions.
+  """
+
+  global_shape : list[int | None] = [None] * len(local_shape)
+  for i, local_dim in enumerate(local_shape):
+    try:
+      _, shard_count = get_process_index_and_count(
+          sharding, i, ndims=len(local_shape)
+      )
+      global_shape[i] = local_dim * shard_count
+    except NonUniformShardingError:
+      global_shape[i] = None
+      continue
+
+  return tuple(global_shape)
+
+
 def num_addressable_indices(
     tensor_sharding: sharding.Sharding,
     dim: int,


Test patch:
diff --git a/jax/_src/test_util.py b/jax/_src/test_util.py
index 8256174f1..193fd71e2 100644
--- a/jax/_src/test_util.py
+++ b/jax/_src/test_util.py
@@ -11,9 +11,11 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
+# pyformat: disable
 from __future__ import annotations
 
-from collections.abc import Generator, Iterable, Mapping, Sequence
+from collections.abc import Generator, Iterable, Sequence
 from contextlib import ExitStack, contextmanager
 import datetime
 import functools
@@ -1961,101 +1963,3 @@ class numpy_with_mpmath:
       return worker(ctx, scale, exact, reference, value)
     else:
       assert 0  # unreachable
-
-
-def get_process_index_and_count(
-    tensor_sharding: jax.sharding.Sharding,
-    dim: int,
-    global_shape: tuple[int, ...],
-) -> tuple[int, int]:
-  """Returns current process index and total count for the given dimension.
-
-  This function facilitates mapping of process-level data to individual
-  devices. Each process can use its index to obtain the data corresponding
-  to that index. If process level data is sharded on multiple dimensions
-  this function can be used to build the cross product of indices in
-  each sharded axis. Processes that need to load the same data will have
-  the same index. For shardings whose per-process data is not distributed
-  on a grid, the number of distinct shards will be such that it is possible to
-  build the target shape while maintaining a "cube" shape of local-process data.
-
-  For example, in case of 4 hosts with sharding distributed like so:
-
-  1234
-  2143
-
-  For dim 0 (rows): all processes need to access all rows, so we return (0, 1)
-  For dim 1 (cols):
-     process 1 and 2 returns index 0 out of 2 (need cols 0 and 1),
-     process 3 and 4 returns index 1 out of 2 (need cols 2 and 3).
-
-  On the other hand, for a sharding like:
-
-  1212
-  3434
-
-  Dim 0 (rows): process 1 and 2 returns (0, 2), process 3 and 4 returns (1, 2)
-  Dim 1 (cols): process 1 and 3 returns (0, 2), process 2 and 4 returns (1, 2)
-
-  Note: This function requires sharding to be process uniform in dimension `dim`:
-   each process has the same number of addressable indices in that
-  dimension and all index sets across processes are either disjoint or the same.
-
-  For sharding to be process uniform the addressable shards doesn't need to
-  form contiguous subtensor, or even a sparse grid  and  in case of
-  interleaved high-dimensional tensor it is possible for sharding to be
-  process uniform only in some dimensions but not others.
-
-  For example:
-    1111 and 12 and 1212 and 1212
-    2222     21     2121     1212
-
-  are all sharding uniform, in both dimensions. However
-
-    1122
-    2121
-    1121
-    1222
-
-  is uniform in dimension 0 (both hosts access all rows), but
-  is not uniform in dimension 1 (host 1 accesses columns: 0, 1, and 3),
-  while host 2 accesses (0, 1, 2, 3).
-
-  Returns:
-    A tuple of (index, num_distinct_shards) for the given dimension.
-    It is guaranteed that `index` will cover 0 to `num_distinct_shards - 1`,
-    across all processes.
-
-  Raises:
-    ValueError if the sharding is not process uniform in dimension `dim`.
-  """
-  # TODO(sandler, yashkatariya): Consider making this function public.
-
-  if tensor_sharding.is_fully_addressable or tensor_sharding.is_fully_replicated:
-    return (0, 1)
-  # NB: For most types of shardings, global_shape is a superfluous argument
-  # and could be replaced by [d, d, ...., d, d], where d is the number of
-  # devices.
-  device_map: Mapping[jax.sharding.Device, jax.sharding.Index] = (
-      tensor_sharding.devices_indices_map(global_shape)
-  )
-
-  global_slice = {k: v[dim] for k, v in device_map.items()}
-  process_map: dict[int, set[tuple[int, int]]] = {}
-  all_slices = set()
-
-  current_pid = next(iter(tensor_sharding.addressable_devices)).process_index
-  for d, v in global_slice.items():
-    key = (v.start, v.stop)
-    process_map.setdefault(d.process_index, set()).add(key)
-    all_slices.add(key)
-  addressable = frozenset(process_map[current_pid])
-  slices_per_process = len(addressable)
-  if any(len(x) != slices_per_process for x in process_map.values()):
-    raise ValueError(f'{tensor_sharding=} is non-uniform on {dim=}')
-  unique_processes = list({frozenset(x) for x in process_map.values()})
-
-  # After removing duplicate processes each slide should appear exactly once.
-  if sum(len(h) for h in unique_processes) != len(all_slices):
-    raise ValueError(f'{tensor_sharding=} is non-uniform on {dim=}')
-  return (unique_processes.index(addressable), len(unique_processes))
diff --git a/tests/mesh_utils_test.py b/tests/mesh_utils_test.py
index ce0778129..1c8893f36 100644
--- a/tests/mesh_utils_test.py
+++ b/tests/mesh_utils_test.py
@@ -21,13 +21,23 @@ import dataclasses
 from absl import logging
 from absl.testing import absltest
 from absl.testing import parameterized
+from jax._src import mesh as mesh_lib
 from jax._src import test_util
+from jax._src.sharding_impls import NamedSharding, PartitionSpec, local_to_global_shape
 from jax.experimental import mesh_utils
 from jax.sharding import Mesh  # pylint: disable=g-importing-member
 import numpy as np
 
 # pyformat: disable
 
+
+@dataclasses.dataclass(frozen=True)
+class MockClient:
+  """Mock client for testing, everything is done as process index 0."""
+  def process_index(self) -> int:
+    return 0
+
+
 @dataclasses.dataclass(frozen=True)
 class MockTpuDevice:
   """Mock TPU device for testing."""
@@ -38,6 +48,7 @@ class MockTpuDevice:
   coords: Sequence[int]
   core_on_chip: int
   slice_index: int = 0
+  client: MockClient = dataclasses.field(default_factory=MockClient)
 
 
 def mock_tpu_devices(x, y, z, dev_kind, one_device_per_chip, num_slices=1,
@@ -207,6 +218,41 @@ class MeshUtilsTest(test_util.JaxTestCase):
         )
     self.assertArraysEqual(assignment, expected_assignment_matrix)
 
+  @parameterized.named_parameters(
+      ('2x2x1', mock_2x2x1_devices,),
+      ('2x2x4', mock_2x2x4_devices, ),
+      ('4x4x4', mock_4x4x4_devices,),
+      ('4x4x8', mock_4x4x8_devices,),
+      ('4x8x8', mock_4x8x8_devices, ),
+      ('8x8', mock_8x8_devices),
+  )
+  def test_create_device_mesh_has_computable_global_shape(self, devices):
+    def factorize(n, max_factors=3):
+      if max_factors == 1 or n == 1:
+        yield (n, ) * max_factors
+        return
+      for i in range(2, n+1):
+        if n % i == 0:
+          for remaining in factorize(n // i, max_factors=max_factors - 1):
+            yield (i, *remaining)
+    jax_devices = devices(True)
+    for mesh_shape in factorize(len(jax_devices), max_factors=3):
+      mesh = mesh_utils.create_device_mesh(mesh_shape, devices=jax_devices,
+                                           allow_split_physical_axes=True)
+      mesh = mesh_lib.Mesh(mesh, ('a', 'b', 'c'))
+      sharding = NamedSharding(mesh, PartitionSpec('a', 'b', 'c'))
+      computed_global_shape = local_to_global_shape(sharding, (1, 1, 1))
+      self.assertFalse(
+          np.any([x is None for x in computed_global_shape]),
+          f'{mesh_shape=}, {computed_global_shape=} is not uniform')
+
+      sharding = NamedSharding(mesh, PartitionSpec(('a', 'c',), 'b'))
+      computed_global_shape = local_to_global_shape(sharding, (1, 1, 1))
+      self.assertFalse(
+          np.any([x is None for x in computed_global_shape]),
+          f'{mesh_shape=}, {computed_global_shape=} is not uniform')
+
+
   @parameterized.named_parameters(
       ('2x2x1', mock_2x2x1_devices, [1, 1, 4], [(), (), (0, 1, 2)]),
       ('2x2x4', mock_2x2x4_devices, [1, 4, 4], [(), (2,), (0, 1)]),