Repository: google/jax
Pull request #22539 (2024-07-20 00:01:57)

Fix the efficient reshard path in device_put when you want to go from 1 mesh to another with different device assignments.
Fix the efficient reshard path in device_put when you want to go from 1 mesh to another with different device assignments.

The old code lead to the wrong answer as shown in the test added in this PR.



Patch:
diff --git a/jax/_src/array.py b/jax/_src/array.py
index 0485bbe2c..796d313ae 100644
--- a/jax/_src/array.py
+++ b/jax/_src/array.py
@@ -32,7 +32,6 @@ from jax._src import dispatch
 from jax._src import dtypes
 from jax._src import errors
 from jax._src import profiler
-from jax._src import tree_util
 from jax._src import xla_bridge
 from jax._src.interpreters import mlir
 from jax._src.interpreters import pxla
@@ -1069,22 +1068,18 @@ def shard_device_array(x, devices, indices, sharding):
   aval = api_util.shaped_abstractify(x)
   return pxla.batched_device_put(aval, sharding, shards, devices)
 
-def _hashable_index(idx):
-  return tree_util.tree_map(
-      lambda x: (x.start, x.stop) if type(x) == slice else x, idx)
-
 
 def shard_sharded_device_array_slow_path(x, devices, indices, sharding):
   candidates = defaultdict(list)
   bufs = [buf.data for buf in x.addressable_shards]
   arr_indices = tuple(x.sharding.devices_indices_map(x.shape).values())
   for buf, idx in safe_zip(bufs, arr_indices):
-    candidates[_hashable_index(idx)].append(buf)
+    candidates[hashed_index(idx)].append(buf)
 
   bufs = []
   for idx, device in safe_zip(indices, devices):
     # Look up all buffers that contain the correct slice of the logical array.
-    candidates_list = candidates[_hashable_index(idx)]
+    candidates_list = candidates[hashed_index(idx)]
     if not candidates_list:
       # This array isn't sharded correctly. Reshard it via host roundtrip.
       # TODO(skye): more efficient reshard?
diff --git a/jax/_src/dispatch.py b/jax/_src/dispatch.py
index 57a70b3ff..6ed3e0924 100644
--- a/jax/_src/dispatch.py
+++ b/jax/_src/dispatch.py
@@ -16,6 +16,7 @@
 from __future__ import annotations
 
 import atexit
+import collections
 from collections.abc import Callable, Iterator, Sequence
 import contextlib
 import dataclasses
@@ -363,13 +364,38 @@ def _different_device_order_reshard(x, target_sharding):
         old_hlo_sharding.tile_assignment_devices(), permute_order
     )
     new_hlo_sharding = xc.HloSharding.from_proto(new_op_sharding)
+    # TODO(yashkatariya): Enable this when HloSharding conversion is fixed in
+    # XLA.
+    # assert (new_op_sharding.tile_assignment_dimensions
+    #         == new_hlo_sharding.tile_assignment_dimensions())
+    # assert (new_op_sharding.tile_assignment_devices
+    #         == new_hlo_sharding.tile_assignment_devices())
+
+  new_sharding = GSPMDSharding(
+      target_sharding._device_assignment, new_hlo_sharding,
+      memory_kind=target_sharding.memory_kind)
+
+  old_device_to_index_buffer = collections.defaultdict()
+  old_index_to_buffer = collections.defaultdict()
+  for s in x.addressable_shards:
+    old_index_to_buffer[array.hashed_index(s.index)] = s.data
+    old_device_to_index_buffer[s.device] = (s.index, s.data)
+
+  new_arrays = []
+  for new_d, new_index in new_sharding.addressable_devices_indices_map(x.shape).items():
+    old_index, old_buf = old_device_to_index_buffer[new_d]
+    if old_index == new_index:
+      assert array._get_device(old_buf) == new_d, (
+          array._get_device(old_buf), new_d)
+      new_arrays.append(old_buf)
+    else:
+      old_buf = old_index_to_buffer[array.hashed_index(new_index)]
+      new_arrays.append(
+          pxla.batched_device_put(old_buf.aval, SingleDeviceSharding(new_d),
+                                  [old_buf], [new_d]))
 
-  new_x = array.make_array_from_single_device_arrays(
-      x.shape,
-      GSPMDSharding(target_sharding._device_assignment, new_hlo_sharding,
-                    memory_kind=target_sharding.memory_kind),
-      x._arrays,
-  )
+  new_x = array.ArrayImpl(
+      x.aval, new_sharding, new_arrays, committed=True, _skip_checks=True)
 
   _orig_get_and_check_device_assignment = pxla._get_and_check_device_assignment.fn
   pxla._get_and_check_device_assignment.fn = partial(


Test patch:
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index 3b56a4f70..bc37a3798 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -4241,6 +4241,35 @@ class ArrayPjitTest(jtu.JaxTestCase):
     self.assertArraysEqual(out, np_inp)
     self.assertEqual(out.sharding, s2)
 
+  @parameterized.named_parameters(
+      ("8_2", (8, 2)),
+      ("8_384", (8, 384)),
+  )
+  def test_device_put_efficient_reshard_complex_mesh(self, shape):
+    if jax.device_count() < 8:
+      self.skipTest('Requires >= 8 devices')
+
+    dev = jax.devices()
+    mesh1 = jax.sharding.Mesh(
+        np.asarray(dev).reshape([1, 2, 2, 2]),
+        ('replica', 'data', 'seq', 'model'))
+    mesh2 = jax.sharding.Mesh(
+        np.asarray(jax.devices())
+        .reshape([1, 1, 2, 2, 2, 1])
+        .swapaxes(2, 3)
+        .reshape([1, 1, 4, 2, 1]),
+        ('replica', 'data', 'seq', 'model_q', 'model_kv'))
+
+    np_inp = jnp.arange(math.prod(shape)).reshape(shape)
+    s1 = NamedSharding(mesh1, P('model'))
+    s2 = NamedSharding(mesh2, P())
+
+    x_s1 = jax.device_put(np_inp, s1)
+
+    out = jax.device_put(x_s1, s2)
+    self.assertArraysEqual(out, np_inp)
+    self.assertEqual(out.sharding, s2)
+
   def test_convert_element_type_sharding(self):
     mesh = jtu.create_global_mesh((2, 2), ('x', 'y'))
     s = NamedSharding(mesh, P('x', 'y'))