Add multi-host support to Safetensors loading in Orbax.

checkpoint/orbax/checkpoint/experimental/v1/_src/context/context.py

@@ -119,6 +119,7 @@ def __init__(
       checkpoint_layout: options_lib.CheckpointLayout | None = None,
       deletion_options: options_lib.DeletionOptions | None = None,
       memory_options: options_lib.MemoryOptions | None = None,
+      safetensors_options: options_lib.SafetensorsOptions | None = None,
   ):
     self._pytree_options = pytree_options or (
         context.pytree_options if context else options_lib.PyTreeOptions()
@@ -156,6 +157,11 @@ def __init__(
     self._memory_options = memory_options or (
         context.memory_options if context else options_lib.MemoryOptions()
     )
+    self._safetensors_options = safetensors_options or (
+        context.safetensors_options
+        if context
+        else options_lib.SafetensorsOptions()
+    )
 
   @property
   def pytree_options(self) -> options_lib.PyTreeOptions:
@@ -197,6 +203,10 @@ def deletion_options(self) -> options_lib.DeletionOptions:
   def memory_options(self) -> options_lib.MemoryOptions:
     return self._memory_options
 
+  @property
+  def safetensors_options(self) -> options_lib.SafetensorsOptions:
+    return self._safetensors_options
+
   def operation_id(self) -> str:
     return synchronization.OperationIdGenerator.get_current_operation_id()
 

checkpoint/orbax/checkpoint/experimental/v1/_src/context/options.py

@@ -570,6 +570,19 @@ class MemoryOptions:
   is_prioritized_key_fn: serialization_types.IsPrioritizedKeyFn | None = None
 
 
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class SafetensorsOptions:
+  """Options for configuring Safetensors loading.
+
+  Attributes:
+    ignore_load_sharding: If True, skips sharding of the tensors across
+      hosts/devices during load. Whole tensors will be present on each host,
+      allowing for efficient conversion.
+  """
+
+  ignore_load_sharding: bool = False
+
+
 class CheckpointLayout(enum.Enum):
   """The layout of the checkpoint.