#p2p Support CPU memory ops in multislice functions

PiperOrigin-RevId: 869791115

Author: Orbax Authors

checkpoint/orbax/checkpoint/_src/multihost/multislice.py

@@ -15,6 +15,8 @@
 """Multislice utilities."""
 
 import functools
+import math
+import os
 from typing import Any, Optional, Set, Union
 
 from absl import logging
@@ -139,6 +141,10 @@ def in_replica(
 def get_device_memory() -> int:
   """Returns HBM capacity of the device on which the code is running(in bytes)."""
   device = jax.devices()[0]
+  if device.platform == 'cpu':
+    page_size = os.sysconf('SC_PAGE_SIZE')
+    phys_pages = os.sysconf('SC_PHYS_PAGES')
+    return int(page_size * phys_pages)
   if device.platform not in ('tpu', 'gpu'):
     raise ValueError('Only select TPU and GPU devices are supported.')
   hbm_memory = {
@@ -163,6 +169,9 @@ def get_device_memory() -> int:
 
 def get_leaf_memory_per_device(arr: jax.Array) -> int:
   """Returns the memory usage of a sharded array per device (in bytes)."""
+  if not arr.addressable_shards:
+    shard_shape = arr.sharding.shard_shape(arr.shape)
+    return math.prod(shard_shape) * arr.dtype.itemsize
   shard = arr.addressable_shards[0]
   return shard.data.size * shard.data.itemsize