Fix pinned_host loading by ensuring memory_kind is carried through when creating in-memory buffers.

PiperOrigin-RevId: 738822867

Author: cpgaffney1

checkpoint/CHANGELOG.md

@@ -7,11 +7,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## [0.11.10] - 2025-03-20
+
 ### Added
 
 - Add `fallback_sharding` option to `StandardRestoreArgs` to support restoring
 on different topologies easily.
 
+### Fixed
+
+- Fix pinned_host loading by ensuring `memory_kind` is carried through when creating in-memory buffers.
+
 ## [0.11.9] - 2025-03-17
 
 ### Added

checkpoint/orbax/checkpoint/_src/handlers/pytree_checkpoint_handler_test_utils.py

@@ -2378,3 +2378,24 @@ def test_save_restore_random_keys(self, use_ocdbt: bool):
             self.directory, args=PyTreeRestoreArgs(pytree, restore_args)
         )
         test_utils.assert_tree_equal(self, pytree, restored)
+
+    def test_pinned_host_loading(self):
+      if multihost.is_pathways_backend():
+        # TODO(b/404915487): Reenable when possible.
+        self.skipTest('Disabled due to b/404915487.')
+      pytree = dict(arr=np.ones((1024, 512)))
+      self.handler.save(self.directory, args=PyTreeSaveArgs(pytree))
+
+      mesh = jax.sharding.Mesh(
+          np.asarray(jax.devices()).reshape((1, len(jax.devices()))), ('x', 'y')
+      )
+      sharding = jax.sharding.NamedSharding(
+          mesh, jax.sharding.PartitionSpec('x', 'y')
+      ).with_memory_kind('pinned_host')
+
+      restore_args = dict(arr=ArrayRestoreArgs(sharding=sharding))
+      restored = self.handler.restore(
+          self.directory, args=PyTreeRestoreArgs(restore_args=restore_args)
+      )
+      expected = dict(arr=jax.device_put(np.ones((1024, 512)), sharding))
+      self.validate_restore(expected, restored)

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

@@ -401,6 +401,7 @@ async def _read_array_index_and_device_put(
     byte_limiter: ByteLimiter,
     strict: bool,
     dll: Optional[layout.DeviceLocalLayout],
+    memory_kind: Optional[str],
 ) -> list[jax.Array]:
   """Callback that reads an array index and places on the devices."""
   for sl in index:
@@ -450,11 +451,10 @@ async def _read_array_index_and_device_put(
           f' TensorStore details: {t.spec}.'
       ) from e
     for device in devices:
-      result.append(
-          jax.device_put(
-              shard, Layout(dll, jax.sharding.SingleDeviceSharding(device))
-          )
+      sharding = jax.sharding.SingleDeviceSharding(
+          device, memory_kind=memory_kind
       )
+      result.append(jax.device_put(shard, Layout(dll, sharding)))
   return result
 
 
@@ -496,6 +496,7 @@ async def read_and_create_array(
           byte_limiter=byte_limiter,
           strict=strict,
           dll=dll,
+          memory_kind=sharding.memory_kind,
       )
       for idx, devices in local_indices_devices_map.items()
   ]

checkpoint/orbax/checkpoint/experimental/v1/_src/handlers/pytree_handler_test_base.py

@@ -1879,3 +1879,26 @@ def test_save_restore_random_keys(self, use_ocdbt: bool):
       ) as load_handler:
         restored = load_handler.load(self.directory)
         test_utils.assert_tree_equal(self, pytree, restored)
+
+    def test_pinned_host_loading(self):
+      if multihost.is_pathways_backend():
+        # TODO(b/404915487): Reenable when possible.
+        self.skipTest('Disabled due to b/404915487.')
+      pytree = dict(arr=np.ones((1024, 512)))
+      self.handler.save(self.directory, pytree)
+
+      mesh = jax.sharding.Mesh(
+          np.asarray(jax.devices()).reshape((1, len(jax.devices()))), ('x', 'y')
+      )
+      sharding = jax.sharding.NamedSharding(
+          mesh, jax.sharding.PartitionSpec('x', 'y')
+      ).with_memory_kind('pinned_host')
+
+      abstract_pytree = dict(
+          arr=jax.ShapeDtypeStruct(
+              pytree['arr'].shape, pytree['arr'].dtype, sharding=sharding
+          )
+      )
+      restored = self.handler.load(self.directory, abstract_pytree)
+      expected = dict(arr=jax.device_put(np.ones((1024, 512)), sharding))
+      test_utils.assert_tree_equal(self, expected, restored)

checkpoint/orbax/checkpoint/version.py

@@ -17,7 +17,7 @@
 # A new PyPI release will be pushed everytime `__version__` is increased.
 # Also modify version and date in CHANGELOG.
 # LINT.IfChange
-__version__ = '0.11.9'
+__version__ = '0.11.10'
 # LINT.ThenChange(//depot//orbax/checkpoint/CHANGELOG.md)