Fix slice_in_dim device incompatibility with global mesh

PiperOrigin-RevId: 831039670

Author: JustinPan-goog

checkpoint/orbax/checkpoint/_src/serialization/replica_slices.py

@@ -46,6 +46,7 @@ class SliceArgs:
   Intended to be passed to `ReplicaSlice` in order to select a slice of
   the `unsliced_data` array.
   """
+
   start_index: int
   limit_index: int
   axis: int
@@ -81,15 +82,28 @@ def is_on_host(self):
     return isinstance(self.unsliced_data, np.ndarray)
 
   def data(self):
+    """Returns the sliced data for this replica.
+
+    If `slice_args` is None, the entire `unsliced_data` is returned. Otherwise,
+    a slice of `unsliced_data` is returned based on `slice_args`.
+    """
     if self.slice_args is None:
       return self.unsliced_data
     else:
-      return jax.lax.slice_in_dim(
-          self.unsliced_data,
-          start_index=self.slice_args.start_index,
-          limit_index=self.slice_args.limit_index,
-          axis=self.slice_args.axis,
+      # If a global mesh is set, slice_in_dim can fail with incompatible device
+      # errors. To avoid this, we temporarily set a mesh constructed from
+      # array's devices.
+      mesh = jax.sharding.Mesh(
+          np.array(list(self.unsliced_data.sharding.device_set)), ('data',)
       )
+      with jax.sharding.set_mesh(mesh):
+        sliced_data = jax.lax.slice_in_dim(
+            self.unsliced_data,
+            start_index=self.slice_args.start_index,
+            limit_index=self.slice_args.limit_index,
+            axis=self.slice_args.axis,
+        )
+      return sliced_data
 
 
 @dataclasses.dataclass(frozen=True)
@@ -421,10 +435,7 @@ def use_pinned_host_transfer(device: jax.Device):
     has_pinned_host = any(
         m.kind == 'pinned_host' for m in device.addressable_memories()
     )
-    return (
-        enable_pinned_host_transfer
-        and has_pinned_host
-    )
+    return enable_pinned_host_transfer and has_pinned_host
 
   def async_transfer_slice(
       rslice: ReplicaSlice,

checkpoint/orbax/checkpoint/checkpoint_manager_test.py

@@ -3907,6 +3907,28 @@ def test_partial_restore_with_omission(self):
         )
         test_utils.assert_tree_equal(self, expected, restored)
 
+  @parameterized.parameters(True, False)
+  def test_save_with_global_mesh(self, use_same_mesh: bool):
+    if use_same_mesh:
+      devices = np.asarray(jax.devices())
+      axis_names = ('x',)
+    else:
+      if multihost.is_pathways_backend():
+        self.skipTest('Not applicable to Pathways.')
+      devices = np.asarray(jax.devices()[:4])
+      axis_names = ('x',)
+    mesh = jax.sharding.Mesh(devices, axis_names)
+    jax.sharding.set_mesh(mesh)
+
+    with CheckpointManager(
+        self.directory,
+        item_names=('params',),
+    ) as manager:
+      self.assertTrue(self.save_params(0, manager, self.pytree))
+      self.wait_if_async(manager)
+      restored = self.restore_params(0, manager)
+      test_utils.assert_tree_equal(self, self.pytree, restored)
+
 
 if __name__ == '__main__':
   multiprocess_test.main()

checkpoint/orbax/checkpoint/experimental/v1/_src/testing/save_load_test_base.py

@@ -926,6 +926,21 @@ def test_partial_restore_omission(self):
 
       test_utils.assert_tree_equal(self, expected, loaded)
 
+    @parameterized.parameters(True, False)
+    def test_save_with_global_mesh(self, use_same_mesh: bool):
+      if use_same_mesh:
+        devices = np.asarray(jax.devices())
+        axis_names = ('x',)
+      else:
+        devices = np.asarray(jax.devices()[:4])
+        axis_names = ('x',)
+      mesh = jax.sharding.Mesh(devices, axis_names)
+      jax.sharding.set_mesh(mesh)
+
+      ocp.save_pytree(self.directory, self.pytree)
+      loaded = ocp.load_pytree(self.directory, self.abstract_pytree)
+      test_utils.assert_tree_equal(self, self.pytree, loaded)
+
     @parameterized.parameters((3,), (8,))
     def test_primary_host_background_error(self, timeout):
       def _assert_false(*args, **kwargs):