Internal Change

FUTURE_COPYBARA_INTEGRATE_REVIEW=#3455 from mridul-sahu:feature/safetensors-sharding-driven-load 739128d
PiperOrigin-RevId: 947822041

Author: mridul-sahu

checkpoint/CHANGELOG.md

@@ -13,6 +13,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Support explicit devices for Pathways MTC init
 - Avoid stale Pathways MTC worker dummy handles
 
+### Changed
+
+- #v1 Drive the safetensors load by the target sharding: each process reads
+only the byte ranges its own devices need (coalesced under
+`SafetensorsOptions.max_over_read_ratio`, fetched as concurrent ranged
+reads of `SafetensorsOptions.read_chunk_bytes` each, bounded by
+`MemoryOptions.read_concurrent_bytes`), replacing the transient-array
+resharding load. Supports arbitrary shardings, including inner-dimension
+partitions.
+
 ## [0.12.1] - 2026-06-24
 
 ### Added

checkpoint/orbax/checkpoint/_src/arrays/types.py

@@ -31,6 +31,7 @@
 NdSlice = tuple[slice, ...] | type(Ellipsis)
 
 Index = tuple[slice, ...]
+IndexBounds = tuple[tuple[int, int], ...]
 
 
 @dataclasses.dataclass(frozen=True)

checkpoint/orbax/checkpoint/_src/testing/benchmarks/run_benchmarks.py

@@ -20,12 +20,16 @@
 using standard environment variables like JAX_COORDINATOR_ADDRESS,
 JAX_PROCESS_ID and JAX_NUM_PROCESSES.
 """
-# pylint: disable=g-statement-before-imports,g-import-not-at-top
+# pylint: disable=g-statement-before-imports,g-import-not-at-top,g-bad-import-order
 
-try:  # SimDevice import must occur before JAX.
-  import simdevice  # pylint: disable=unused-import
-except (ImportError, FileNotFoundError):
-  pass
+import sys
+
+# SimDevice import must occur before JAX only if simulated run is requested.
+if any('simdevice' in arg for arg in sys.argv):
+  try:
+    import simdevice  # pylint: disable=unused-import
+  except (ImportError, FileNotFoundError):
+    pass
 
 import os
 

checkpoint/orbax/checkpoint/_src/testing/benchmarks/safetensors/configs/llama3-405b.yaml

@@ -8,8 +8,11 @@
 #     of byte runs to bin-pack);
 #   * per-host shard sizes are tiny relative to file sizes (4 GB
 #     shards out of 100+ GB files);
-#   * concurrent reads cap (`_MAX_CONCURRENT_READS`) is the active
-#     bottleneck, not bandwidth.
+#   * the byte budget (`max_in_flight_bytes`, default 2 GiB) governs
+#     concurrency by total bytes — so on this tier each read is small
+#     (~4 GB per host, easily within budget) and many reads run
+#     concurrently. Raise the budget on RAM-rich hosts to absorb the
+#     larger coalesced blocks cross-tensor coalescing produces.
 #
 # Llama 3 405B is gated; see safetensor_load_llama3-8b.yaml.
 # GCS pre-staging is *strongly* recommended -- 810 GB download per VM

checkpoint/orbax/checkpoint/_src/testing/benchmarks/safetensors/configs/mistral7b.yaml

@@ -12,7 +12,7 @@
 # bf16 weights at fsdp=8 -> ~1.75 GB per core, comfortable on any v5p/v5e.
 
 suite_name: "safetensors_load_mistral7b"
-num_repeats: 1
+num_repeats: 2
 
 mesh_config:
   mesh_axes: ["fsdp"]
@@ -23,5 +23,7 @@ benchmarks:
     options:
       checkpoint_path: "gs://orbax-benchmarks/st_fixtures/mistral7b/"
       sharding_config_path: "gs://orbax-benchmarks/st_fixtures/sharding/mistral7b_fsdp8.json"
-      enable_trace: true
-      # Capture-only.
+      enable_trace: false
+      enable_xprof: true
+      safetensors_max_over_read_ratio: [1.0]
+      safetensors_read_chunk_mb: [256]

checkpoint/orbax/checkpoint/_src/testing/benchmarks/safetensors/configs/mistral7b_tp.yaml

@@ -31,7 +31,7 @@
 #     new loader's reads should be bounded at 2x ideal, not K x.
 
 suite_name: "safetensor_load_mistral7b_tp"
-num_repeats: 4
+num_repeats: 2
 
 mesh_config:
   mesh_axes: ["tp"]
@@ -41,5 +41,9 @@ benchmarks:
   - generator: "orbax.checkpoint._src.testing.benchmarks.safetensors.load_benchmark.SafetensorLoadBenchmark"
     options:
       checkpoint_path: "gs://orbax-benchmarks/st_fixtures/mistral7b/"
-      sharding_config_path: "gs://orbax-benchmarks/st_fixtures/sharding/mistral7b_tp8.json"
-      enable_trace: true
+      sharding_config_path: "gs://orbax-benchmarks/st_fixtures/sharding/mistral7b_fsdp8.json"
+      enable_trace: false
+      enable_xprof: true
+      # The sweep. Each value becomes a separate run in the dashboard;
+      # the parallel-coordinates view colors them so the curve is visible.
+      safetensors_max_over_read_ratio: [2.0]

checkpoint/orbax/checkpoint/_src/testing/benchmarks/safetensors/load_benchmark.py

@@ -28,6 +28,7 @@
 
 from __future__ import annotations
 
+import contextlib
 import dataclasses
 import json
 import pprint
@@ -44,14 +45,16 @@
 from orbax.checkpoint._src.testing.benchmarks.core import pytree_utils
 
 
+
 @dataclasses.dataclass(frozen=True)
 class SafetensorLoadBenchmarkOptions(benchmarks_core.BenchmarkOptions):
   """Configuration for loading a HuggingFace-format safetensors checkpoint.
 
   Self-contained: builds its own `ocp.Context()` with the SAFETENSORS layout
   pinned from the load-side tuning knobs, so a single run can sweep e.g.
-  `use_load_and_broadcast` or `restore_concurrent_gb` for A/B comparison. Each
-  knob may be a single value or a list to expand into a parameter sweep.
+  `use_load_and_broadcast`, `restore_concurrent_gb`, or
+  `safetensors_max_over_read_ratio` for A/B comparison. Each knob may be a
+  single value or a list to expand into a parameter sweep.
 
   Attributes:
     checkpoint_path: Path (local or `gs://`) to a directory containing one or
@@ -74,6 +77,17 @@ class SafetensorLoadBenchmarkOptions(benchmarks_core.BenchmarkOptions):
       others instead of every replica reading from storage.
     restore_concurrent_gb: Cap on concurrent read bytes (in GiB); `None` leaves
       the Orbax default.
+    safetensors_max_over_read_ratio: Override for
+      `SafetensorsOptions.max_over_read_ratio`. Bounds per-host cluster egress
+      on strided shardings (inner-dim TP, HSDP). Sweep as `[2.0, 3.0, 4.0]` on a
+      TP variant to measure the read-count vs over-read tradeoff. `None`
+      (default) uses the loader's default (currently 2.0). In-flight read bytes
+      are bounded by `restore_concurrent_gb` (shared with the rest of restore).
+    safetensors_read_chunk_mb: Override for
+      `SafetensorsOptions.read_chunk_bytes`, in MiB. Coalesced blocks are split
+      at this size into concurrent ranged reads; sweep as `[64, 128, 512]` to
+      measure the request-count vs read-parallelism tradeoff on a given storage
+      backend. `None` (default) uses the loader's default (currently 128 MiB).
     metric_tracemalloc_enabled: Whether to capture the tracemalloc metric
       (opt-in because its per-allocation snapshots are expensive).
   """
@@ -84,7 +98,10 @@ class SafetensorLoadBenchmarkOptions(benchmarks_core.BenchmarkOptions):
   reference_digests_path: str | None = None
   use_load_and_broadcast: bool | list[bool] = False
   restore_concurrent_gb: int | None | list[int | None] = None
+  safetensors_max_over_read_ratio: float | None | list[float | None] = None
+  safetensors_read_chunk_mb: int | None | list[int | None] = None
   metric_tracemalloc_enabled: bool = False
+  enable_xprof: bool = False
 
   def is_valid(self) -> bool:
     if self.checkpoint_path is None:
@@ -94,7 +111,15 @@ def is_valid(self) -> bool:
   @property
   def context(self) -> ocp.Context:
     ctx = ocp.Context(
-        checkpoint_layout=ocp.options.CheckpointLayout.SAFETENSORS
+        checkpoint_layout=ocp.options.CheckpointLayout.SAFETENSORS,
+        safetensors_options=ocp.options.SafetensorsOptions(
+            max_over_read_ratio=self.safetensors_max_over_read_ratio,
+            read_chunk_bytes=(
+                self.safetensors_read_chunk_mb * 1024**2
+                if self.safetensors_read_chunk_mb is not None
+                else None
+            ),
+        ),
     )
     # TODO(b/519204863): Fix type hint for args like list[T] | T
     ctx.array.loading.use_load_and_broadcast = self.use_load_and_broadcast
@@ -214,12 +239,18 @@ def test_fn(
       else:
         abstract_pytree = _replicated_abstract_state(metadata.metadata)
 
+      xprof_session = contextlib.nullcontext()
+
       if load_trace is not None:
         jax.profiler.start_trace(str(load_trace))
-      with metrics.measure("load", metrics_to_measure):
-        restored_pytree = ocp.load(
-            checkpoint_path, abstract_state=abstract_pytree
-        )
+      with xprof_session as url:
+        if url:
+          logging.info("XProf Session URL: %s", url)
+          print(f"\n[XPROF] Session URL: {url}\n")
+        with metrics.measure("load", metrics_to_measure):
+          restored_pytree = ocp.load(
+              checkpoint_path, abstract_state=abstract_pytree
+          )
       if load_trace is not None:
         jax.profiler.stop_trace()
 

checkpoint/orbax/checkpoint/_src/testing/benchmarks/safetensors/load_benchmark_test.py

@@ -55,6 +55,16 @@ def test_context_pins_safetensors_layout(self):
         ocp.options.CheckpointLayout.SAFETENSORS,
     )
 
+  def test_context_maps_safetensors_knobs(self):
+    opts = slb.SafetensorLoadBenchmarkOptions(
+        checkpoint_path="/tmp/x",
+        safetensors_max_over_read_ratio=3.0,
+        safetensors_read_chunk_mb=64,
+    )
+    ctx = opts.context
+    self.assertEqual(ctx.safetensors_options.max_over_read_ratio, 3.0)
+    self.assertEqual(ctx.safetensors_options.read_chunk_bytes, 64 * 1024**2)
+
 
 class GenerationTest(absltest.TestCase):
 

checkpoint/orbax/checkpoint/_src/testing/benchmarks/safetensors/prepare.py

@@ -190,6 +190,8 @@ def _mirror_to_gcs(local_dir: epath.Path, gcs_uri: str) -> None:
   )
 
 
+
+
 def _write_sharding(
     model_dir: epath.Path, axis_size: int, strategy: str, out: str
 ) -> int:
@@ -232,15 +234,20 @@ def _write_sharding(
 flags.mark_flags_as_mutual_exclusive(["repo", "local_dir"], required=True)
 
 
-def main(argv: list[str]) -> None:
-  del argv  # Unused; inputs come from flags.
-  if not _GCS.value and not _SHARDING_OUT.value:
+def _validate_flags() -> None:
+  is_mirroring = bool(_GCS.value)
+  if not is_mirroring and not _SHARDING_OUT.value:
     raise app.UsageError("nothing to do: pass --gcs and/or --sharding_out.")
   if _SHARDING_OUT.value and _AXIS_SIZE.value is None:
     raise app.UsageError("--sharding_out requires --axis_size.")
 
+
+def main(argv: list[str]) -> None:
+  del argv  # Unused; inputs come from flags.
+  _validate_flags()
+
   def _run(model_dir: epath.Path) -> None:
-    if _GCS.value:
+    if _GCS.value and model_dir != epath.Path(_GCS.value):
       print(f">>> Mirroring {model_dir} -> {_GCS.value}")
       _mirror_to_gcs(model_dir, _GCS.value)
     if _SHARDING_OUT.value:
@@ -251,7 +258,7 @@ def _run(model_dir: epath.Path) -> None:
 
   if _REPO.value:
     if _DOWNLOAD_DIR.value:
-      dest = epath.Path(_DOWNLOAD_DIR.value) / "safetensors-prepare-model"
+      dest = epath.Path(_DOWNLOAD_DIR.value)
     else:
       model_slug = _REPO.value.replace("/", "--")
       dest = (

checkpoint/orbax/checkpoint/_src/testing/benchmarks/safetensors/prepare_test.py

@@ -0,0 +1,61 @@
+# Copyright 2026 The Orbax Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from unittest import mock
+from absl import flags
+from absl.testing import absltest
+from etils import epath
+from orbax.checkpoint._src.testing.benchmarks.safetensors import prepare
+
+
+class PrepareTest(absltest.TestCase):
+
+  @mock.patch("huggingface_hub.snapshot_download")
+  @mock.patch("subprocess.run")
+  def test_main_with_repo_and_gcs(self, mock_run, mock_download):
+    flags.FLAGS.repo = "dummy-org/dummy-repo"
+    flags.FLAGS.gcs = "gs://dummy-bucket/model"
+    flags.FLAGS.local_dir = None
+
+    temp_dir = self.create_tempdir()
+    (epath.Path(temp_dir.full_path) / "file1.safetensors").write_text("dummy")
+    (epath.Path(temp_dir.full_path) / "file2.safetensors").write_text("dummy")
+
+    mock_download.return_value = temp_dir.full_path
+
+    prepare.main([])
+
+    mock_download.assert_called_once_with(
+        repo_id="dummy-org/dummy-repo",
+        local_dir=mock.ANY,
+        allow_patterns=["*.safetensors", "*.safetensors.index.json"],
+    )
+    mock_run.assert_called_with(
+        [
+            "gcloud",
+            "storage",
+            "rsync",
+            temp_dir.full_path,
+            "gs://dummy-bucket/model",
+            "-r",
+        ],
+        check=True,
+    )
+
+
+
+if __name__ == "__main__":
+  # Bypass required flag check during test initialization
+  flags.FLAGS.set_default("repo", "dummy-org/dummy-repo")
+  absltest.main()