diff --git a/checkpoint/CHANGELOG.md b/checkpoint/CHANGELOG.md index 5eb6660a5..e01937149 100644 --- a/checkpoint/CHANGELOG.md +++ b/checkpoint/CHANGELOG.md @@ -12,6 +12,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Reduce Memory Pressure from long held CPU buffers in Pathways MTC - Support explicit devices for Pathways MTC init - Avoid stale Pathways MTC worker dummy handles +- #v1 Add `SafetensorsOptions.broadcast_replicated`: read tensors that are +replicated across processes from storage exactly once (spread across the +replicas) and fill every replica with an on-device broadcast, instead of each +process re-reading the same bytes. ### Changed diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/context/options.py b/checkpoint/orbax/checkpoint/experimental/v1/_src/context/options.py index 5cc193a92..04b526286 100644 --- a/checkpoint/orbax/checkpoint/experimental/v1/_src/context/options.py +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/context/options.py @@ -33,7 +33,6 @@ from orbax.checkpoint.experimental.v1._src.serialization import types as serialization_types from orbax.checkpoint.experimental.v1._src.tree import types as tree_types - FROZEN_IDS: contextvars.ContextVar[frozenset[int]] = contextvars.ContextVar( 'orbax_frozen_option_ids', default=frozenset() ) @@ -620,10 +619,20 @@ class SafetensorsOptions(_ActiveContextGuard): (`MemoryOptions.read_concurrent_bytes`), which also caps concurrent requests at `budget / read_chunk_bytes` per host. `None` selects an implementation default (currently 128 MiB). + broadcast_replicated: When True (and running multi-process), the bytes + of tensors that are replicated across processes are read from storage + exactly once -- the read is spread across the replicas -- and every + replica is then filled by an on-device broadcast (an XLA resharding) + instead of each process re-reading the same bytes. Trades accelerator + interconnect traffic for storage egress: enable it when storage + bandwidth or egress cost is the bottleneck (e.g. many hosts loading a + replicated checkpoint from object storage). Applies to leaves with a + `jax.sharding.NamedSharding`; leaves with other shardings load normally. """ max_over_read_ratio: float | None = None read_chunk_bytes: int | None = None + broadcast_replicated: bool = False class CheckpointLayout(enum.Enum): diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/layout/safetensors_layout.py b/checkpoint/orbax/checkpoint/experimental/v1/_src/layout/safetensors_layout.py index 8de433a67..47d3957b1 100644 --- a/checkpoint/orbax/checkpoint/experimental/v1/_src/layout/safetensors_layout.py +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/layout/safetensors_layout.py @@ -39,6 +39,14 @@ one long single-stream read. As each file's reads complete, its tensors are assembled into `jax.Array`s and their host buffers released while other files are still reading. + +`SafetensorsOptions.broadcast_replicated` removes the one case where +sharding-driven reads pull more from storage than the checkpoint's size: +tensors replicated across processes, whose bytes every process would +otherwise re-read. With the flag on, such tensors are read under a +de-duplicated sharding (each byte owned once, spread across the replicas) and +an identity computation with the target `out_shardings` then broadcasts the +shards over the accelerator interconnect. """ import asyncio @@ -49,7 +57,7 @@ import json import math import time -from typing import Any, NamedTuple, cast +from typing import Any, cast, NamedTuple from absl import logging import humanize @@ -114,6 +122,11 @@ # sharding, not the loader, is the thing to change. Over-read from gap-floor # merges alone stays well below this. _OVER_READ_WARN_RATIO = 1.5 +# Per-device output bytes per compiled broadcast under +# `SafetensorsOptions.broadcast_replicated`. Tensors read once under a +# de-duplicated sharding are resharded to their replicated targets in batches +# no larger than this, bounding the transient device memory of each call. +_BROADCAST_BATCH_BYTES = 1024 * 1024 * 1024 # 1 GiB def _get_dtypes() -> dict[str, Any]: @@ -229,7 +242,10 @@ def _normalize_index( for s in resolved: # pyrefly: ignore[not-iterable] if s.step not in (None, 1): raise ValueError(f"Strided shard index is unsupported: {s}.") - return tuple((int(s.start), int(s.stop)) for s in resolved) # pyrefly: ignore[not-iterable] + return tuple( + (int(s.start), int(s.stop)) + for s in resolved # pyrefly: ignore[not-iterable] + ) def _byte_strides(global_shape: arrays_types.Shape, itemsize: int) -> list[int]: @@ -424,6 +440,117 @@ def _replicated_sharding() -> jax.sharding.Sharding: return jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec()) +def _spec_axes(entry: Any) -> tuple[str, ...]: + """Returns one `PartitionSpec` entry's mesh axis names as a tuple.""" + if entry is None: + return () + return entry if isinstance(entry, tuple) else (entry,) + + +def _dedup_read_sharding( + sharding: jax.sharding.Sharding, global_shape: tuple[int, ...] +) -> jax.sharding.NamedSharding | None: + """Returns a sharding under which each tensor byte is owned once, if any. + + In a `NamedSharding`, replication is exactly the mesh axes the spec does not + use: devices that differ only along those axes own identical index domains, + so under `broadcast_replicated` every replica would re-read the same + bytes from storage. Assigning each unused axis to a tensor dimension splits + those identical domains apart -- the same bytes are then owned (and read) + exactly once, spread across all the replicas' hosts. The loader reads under + the returned sharding and afterwards broadcasts to the true target + (`_broadcast_to_target_shardings`). + + Each unused axis is greedily placed on the first dimension that remains + evenly divisible (JAX rejects uneven partitions). Axes that fit nowhere + stay unplaced: the result then still de-duplicates partially. Returns + `None` when the sharding is not a `NamedSharding`, nothing is replicated, + or no axis can be placed -- callers then read under the target sharding + as usual. + + Args: + sharding: The target sharding for the tensor. + global_shape: The global shape of the tensor. + + Returns: + The de-duplicated read sharding, or `None` if none is applicable. + """ + if not global_shape or not isinstance(sharding, jax.sharding.NamedSharding): + return None + mesh = sharding.mesh + entries = [ + _spec_axes(e) + for e in tuple(sharding.spec) + (None,) * len(global_shape) + ][: len(global_shape)] + used = {axis for entry in entries for axis in entry} + unused = [ + axis + for axis in mesh.axis_names + if axis not in used and mesh.shape[axis] > 1 + ] + if not unused: + return None + partitions = [ + math.prod(mesh.shape[axis] for axis in entry) for entry in entries + ] + placed_any = False + for axis in unused: + for dim, dim_size in enumerate(global_shape): + if dim_size % (partitions[dim] * mesh.shape[axis]) == 0: + entries[dim] += (axis,) + partitions[dim] *= mesh.shape[axis] + placed_any = True + break + if not placed_any: + return None + return jax.sharding.NamedSharding( + mesh, jax.sharding.PartitionSpec(*(e or None for e in entries)) + ) + + +def _broadcast_to_target_shardings( + arrays: dict[str, jax.Array], + targets: dict[str, jax.sharding.Sharding], +) -> None: + """Reshards arrays read under a de-duplicated sharding to their targets. + + An array read under `_dedup_read_sharding` holds each byte on exactly one + device; an identity computation whose `out_shardings` are the true targets + makes XLA broadcast the shards to every replica over the accelerator + interconnect. Tensors are batched so each compiled call's per-device output + stays under a fixed byte bound, keeping transient device memory + predictable regardless of checkpoint size. + + Args: + arrays: Loaded arrays by tensor name; entries named in `targets` are + replaced in place with their broadcast result. + targets: Target sharding by tensor name for every array to broadcast. + """ + batch: list[str] = [] + batch_bytes = 0 + + def flush() -> None: + if not batch: + return + resharded = jax.jit( + lambda xs: xs, out_shardings=[targets[n] for n in batch] + )([arrays[n] for n in batch]) + for n, out in zip(batch, resharded): + arrays[n] = out + batch.clear() + + for name in targets: + arr = arrays[name] + shard_shape = targets[name].shard_shape(arr.shape) + per_device_bytes = math.prod(shard_shape) * arr.dtype.itemsize + if batch and batch_bytes + per_device_bytes > _BROADCAST_BATCH_BYTES: + flush() + batch_bytes = 0 + batch.append(name) + batch_bytes += per_device_bytes + flush() + + class _FileEntry(NamedTuple): """Where to find one tensor in the on-disk checkpoint.""" @@ -646,6 +773,7 @@ def _plan_whole_tensor( def _plan_sharded_tensor( entry: _FileEntry, abstract_leaf: Any, + sharding: jax.sharding.Sharding, ) -> tuple[list[_Read], Callable[[], jax.Array]]: """Plans the shard reads this process's devices need for one tensor. @@ -656,7 +784,10 @@ def _plan_sharded_tensor( Args: entry: The file entry describing the tensor's location and header info. - abstract_leaf: The abstract array (shape, dtype, sharding) to build. + abstract_leaf: The abstract array (shape, dtype) to build. + sharding: The sharding to read and assemble under -- the leaf's target + sharding, or the de-duplicated read sharding when + `broadcast_replicated` is enabled. Returns: The tensor's byte runs, and a builder that assembles the populated shard @@ -676,7 +807,6 @@ def _plan_sharded_tensor( f"Shape mismatch: abstract state has {global_shape}, the checkpoint" f" has {file_shape}." ) - sharding = getattr(abstract_leaf, "sharding", None) or _replicated_sharding() this_process = multihost.process_index() index_to_devices: dict[arrays_types.IndexBounds, list[jax.Device]] = ( @@ -801,6 +931,60 @@ def _warn_if_over_read( ) +def _plan_reads( + file_index: dict[str, _FileEntry], + abstract_state: dict[str, Any] | None, + names: Sequence[str], + broadcast_replicated: bool, +) -> tuple[ + dict[Path, list[_Read]], + dict[Path, list[tuple[str, Callable[[], Any]]]], + dict[str, jax.sharding.Sharding], +]: + """Plans every tensor's byte runs and builders, grouped by file. + + With `broadcast_replicated`, replicated bytes are read from storage + once -- under a de-duplicated sharding spread across the replicas -- and + the replicas are filled afterwards by an on-device broadcast. + Cross-process duplication is the only waste it removes, so it is a no-op + for single-process loads. + + Args: + file_index: Tensor name -> `_FileEntry` for the whole checkpoint. + abstract_state: The abstract leaves driving the load, or `None` to + materialise every tensor in full (single-process only). + names: The tensor names to load, in result order. + broadcast_replicated: `SafetensorsOptions.broadcast_replicated`. + + Returns: + Reads grouped by file, `(name, build)` pairs grouped by file, and the + target shardings for every tensor read under a de-duplicated sharding + (the tensors `_broadcast_to_target_shardings` must fix up afterwards). + """ + dedup_replicated = broadcast_replicated and multihost.process_count() > 1 + reads_by_file: dict[Path, list[_Read]] = collections.defaultdict(list) + builds_by_file: dict[Path, list[tuple[str, Callable[[], Any]]]] = ( + collections.defaultdict(list) + ) + broadcast_targets: dict[str, jax.sharding.Sharding] = {} + for name in names: + entry = file_index[name] + if abstract_state is None: + reads, build = _plan_whole_tensor(entry) + else: + leaf = abstract_state[name] + sharding = getattr(leaf, "sharding", None) or _replicated_sharding() + if dedup_replicated and ( + read_sharding := _dedup_read_sharding(sharding, tuple(leaf.shape)) + ): + broadcast_targets[name] = sharding + sharding = read_sharding + reads, build = _plan_sharded_tensor(entry, leaf, sharding) + reads_by_file[entry.path].extend(reads) + builds_by_file[entry.path].append((name, build)) + return reads_by_file, builds_by_file, broadcast_targets + + class SafetensorsLayout(CheckpointLayout): """Handles checkpoints in the HuggingFace Safetensors format. @@ -1009,18 +1193,9 @@ async def _load( # Phase 1 (sync): compute every byte range this process needs, grouped by # file, and pre-allocate the destination buffers. Each tensor's `build()` # turns its (yet-unfilled) buffers into an array once its file is read. - reads_by_file: dict[Path, list[_Read]] = collections.defaultdict(list) - builds_by_file: dict[Path, list[tuple[str, Callable[[], Any]]]] = ( - collections.defaultdict(list) + reads_by_file, builds_by_file, broadcast_targets = _plan_reads( + file_index, abstract_state, names, opts.broadcast_replicated ) - for name in names: - entry = file_index[name] - if abstract_state is None: - reads, build = _plan_whole_tensor(entry) - else: - reads, build = _plan_sharded_tensor(entry, abstract_state[name]) - reads_by_file[entry.path].extend(reads) - builds_by_file[entry.path].append((name, build)) # Phase 2 (async): every file's chunks in parallel, bounded together by # the shared limiter. The chunk reads fill the pre-allocated buffers in @@ -1047,4 +1222,6 @@ async def load_file(file_path: Path) -> _ReadStats: stats = await asyncio.gather(*map(load_file, list(reads_by_file))) _record_read_stats(stats) _warn_if_over_read(stats, opts.max_over_read_ratio is None) + if broadcast_targets: + _broadcast_to_target_shardings(arrays, broadcast_targets) return {name: arrays[name] for name in names} diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/layout/safetensors_layout_multiprocess_test.py b/checkpoint/orbax/checkpoint/experimental/v1/_src/layout/safetensors_layout_multiprocess_test.py index ea84b45fc..5cf7e47fd 100644 --- a/checkpoint/orbax/checkpoint/experimental/v1/_src/layout/safetensors_layout_multiprocess_test.py +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/layout/safetensors_layout_multiprocess_test.py @@ -17,6 +17,7 @@ import gc import tracemalloc import unittest +from unittest import mock from absl.testing import parameterized from etils import epath @@ -342,5 +343,161 @@ async def test_load_sharded_fails_with_wrong_key_abstract_pytree(self): await test_awaitable +class BroadcastReplicatedTest( + unittest.IsolatedAsyncioTestCase, + parameterized.TestCase, + multiprocess_test.MultiProcessTest, +): + """`SafetensorsOptions.broadcast_replicated` on the 4-process mesh.""" + + def setUp(self): + super().setUp() + self.test_dir = epath.Path( + self.multiprocess_create_tempdir(name="test_dir") + ) + devices = jax.devices() + self.mesh = Mesh(np.array(devices).reshape(4, 2), ("data", "model")) + test_utils.sync_global_processes("BroadcastReplicatedTest.setUp") + + def tearDown(self): + super().tearDown() + test_utils.sync_global_processes("BroadcastReplicatedTest.tearDown") + + def _assert_domains_cover_once(self, sharding, shape): + """Every element owned by exactly one unique index domain (no replicas).""" + domains = { + safetensors_layout._normalize_index(index, shape) + for index in sharding.devices_indices_map(shape).values() + } + covered = sum( + np.prod([stop - start for start, stop in bounds]) for bounds in domains + ) + self.assertEqual(covered, np.prod(shape)) + + def test_dedup_fully_replicated_covers_once(self): + sharding = NamedSharding(self.mesh, PartitionSpec()) + read = safetensors_layout._dedup_read_sharding(sharding, (16, 8)) + self.assertIsNotNone(read) + self._assert_domains_cover_once(read, (16, 8)) + + def test_dedup_partially_replicated_covers_once(self): + sharding = NamedSharding(self.mesh, PartitionSpec(None, "model")) + read = safetensors_layout._dedup_read_sharding(sharding, (16, 8)) + self.assertIsNotNone(read) + self._assert_domains_cover_once(read, (16, 8)) + + def test_dedup_places_axis_on_later_divisible_dim(self): + # dim 0 (size 3) cannot absorb either mesh axis; dim 1 (size 8) takes + # both. JAX requires even partitions, so placement must skip dim 0. + sharding = NamedSharding(self.mesh, PartitionSpec()) + read = safetensors_layout._dedup_read_sharding(sharding, (3, 8)) + self.assertIsNotNone(read) + self._assert_domains_cover_once(read, (3, 8)) + + def test_dedup_fully_sharded_returns_none(self): + sharding = NamedSharding(self.mesh, PartitionSpec("data", "model")) + self.assertIsNone( + safetensors_layout._dedup_read_sharding(sharding, (16, 8)) + ) + + def test_dedup_nothing_placeable_returns_none(self): + # No dimension of a (3, 3) tensor divides evenly by either mesh axis. + sharding = NamedSharding(self.mesh, PartitionSpec()) + self.assertIsNone( + safetensors_layout._dedup_read_sharding(sharding, (3, 3)) + ) + + def test_dedup_partial_placement_still_reduces_replicas(self): + # Only "model" (size 2) fits a (2, 2) tensor; "data" stays unplaced, so + # 2 unique domains remain across 8 devices instead of 1 -- a partial + # de-duplication (4 owners per domain down from 8). + sharding = NamedSharding(self.mesh, PartitionSpec()) + read = safetensors_layout._dedup_read_sharding(sharding, (2, 2)) + self.assertIsNotNone(read) + domains = { + safetensors_layout._normalize_index(index, (2, 2)) + for index in read.devices_indices_map((2, 2)).values() + } + self.assertLen(domains, 2) + + @parameterized.named_parameters( + ("fully_replicated", PartitionSpec(), (16, 512)), + ("partially_replicated", PartitionSpec(None, "model"), (16, 512)), + ("indivisible_leading_dim", PartitionSpec(), (3, 8)), + ) + async def test_load_with_broadcast_matches_direct_device_put( + self, spec, shape + ): + tensor = np.arange(np.prod(shape), dtype=np.float32).reshape(shape) + st_path = self.test_dir / f"{self.id()}.safetensors" + if jax.process_index() == 0: + np_save_file({"w": tensor}, st_path) + test_utils.sync_global_processes(self.id()) + + target = NamedSharding(self.mesh, spec) + abstract_state = { + "w": jax.ShapeDtypeStruct( + shape=shape, dtype=np.float32, sharding=target + ) + } + layout = SafetensorsLayout() + with context_lib.Context( + safetensors_options=options_lib.SafetensorsOptions( + broadcast_replicated=True, + ), + ): + restore_fn = await layout.load(st_path, abstract_state=abstract_state) + restored = (await restore_fn)["w"] + + expected = jax.device_put(tensor, target) + self.assertEqual(restored.sharding, expected.sharding) + test_utils.assert_array_equal(self, expected, restored) + + async def test_broadcast_reads_each_byte_once_per_cluster(self): + # Fully replicated target: without the flag every process reads the + # whole tensor (4x cluster egress); with it, the read spreads so each + # process pulls ~1/4. Each process asserts its own bytes_read metric. + shape = (16, 4096) + tensor = np.arange(np.prod(shape), dtype=np.float32).reshape(shape) + total_bytes = tensor.nbytes + st_path = self.test_dir / f"{self.id()}.safetensors" + if jax.process_index() == 0: + np_save_file({"w": tensor}, st_path) + test_utils.sync_global_processes(self.id()) + + target = NamedSharding(self.mesh, PartitionSpec()) + abstract_state = { + "w": jax.ShapeDtypeStruct( + shape=shape, dtype=np.float32, sharding=target + ) + } + layout = SafetensorsLayout() + + async def bytes_read_for(broadcast_replicated): + with mock.patch.object(jax.monitoring, "record_scalar") as rec: + with context_lib.Context( + safetensors_options=options_lib.SafetensorsOptions( + broadcast_replicated=broadcast_replicated, + ), + ): + restore_fn = await layout.load( + st_path, abstract_state=abstract_state + ) + restored = (await restore_fn)["w"] + test_utils.assert_array_equal( + self, jax.device_put(tensor, target), restored + ) + emitted = {c.args[0]: c.args[1] for c in rec.call_args_list} + return emitted["/jax/orbax/read/safetensors/bytes_read"] + + replicated_bytes = await bytes_read_for(False) + test_utils.sync_global_processes(f"{self.id()}.baseline") + deduped_bytes = await bytes_read_for(True) + + self.assertEqual(replicated_bytes, total_bytes) + # This process's share is 1/4 of the tensor; allow coalescing slack. + self.assertLessEqual(deduped_bytes, 0.3 * total_bytes) + + if __name__ == "__main__": multiprocess_test.main() diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/layout/safetensors_layout_test.py b/checkpoint/orbax/checkpoint/experimental/v1/_src/layout/safetensors_layout_test.py index 0f495ad90..b784dd351 100644 --- a/checkpoint/orbax/checkpoint/experimental/v1/_src/layout/safetensors_layout_test.py +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/layout/safetensors_layout_test.py @@ -458,7 +458,8 @@ def test_normalize_index_open_ended(self): def test_normalize_index_integer(self): self.assertEqual( safetensors_layout._normalize_index( - (2, slice(None, None)), (8, 16) # pyrefly: ignore[bad-argument-type] + (2, slice(None, None)), # pyrefly: ignore[bad-argument-type] + (8, 16), ), # pytype: disable=wrong-arg-types ((2, 3), (0, 16)), ) @@ -969,6 +970,61 @@ async def test_non_positive_read_chunk_bytes_raises(self): with self.assertRaisesRegex(ValueError, 'read_chunk_bytes'): await restore_fn + async def test_broadcast_replicated_is_noop_on_single_process(self): + # A single process already reads each byte exactly once, so the flag + # must not change behavior (and must not attempt any broadcast). + np_save_file({'a': np.arange(8, dtype=np.float32)}, self.file_path) + layout = SafetensorsLayout() + sharding = jax.sharding.NamedSharding( + jax.sharding.Mesh(np.array(jax.devices()), ('x',)), + jax.sharding.PartitionSpec(), + ) + tree = { + 'a': jax.ShapeDtypeStruct( + shape=(8,), dtype=np.float32, sharding=sharding + ) + } + with context_lib.Context( + safetensors_options=options_lib.SafetensorsOptions( + broadcast_replicated=True, + ), + ): + restore_fn = await layout.load(self.file_path, abstract_state=tree) + pytree = await restore_fn + np.testing.assert_allclose(pytree['a'], np.arange(8, dtype=np.float32)) + + +class DedupReadShardingGateTest(absltest.TestCase): + """Cases where `_dedup_read_sharding` must decline (single-device host). + + The placement logic needs a multi-device mesh and is exercised in the + multiprocess suite; here we pin the gates that need no replication. + """ + + def test_scalar_returns_none(self): + sharding = jax.sharding.NamedSharding( + jax.sharding.Mesh(np.array(jax.devices()), ('x',)), + jax.sharding.PartitionSpec(), + ) + self.assertIsNone(safetensors_layout._dedup_read_sharding(sharding, ())) + + def test_non_named_sharding_returns_none(self): + sharding = jax.sharding.SingleDeviceSharding(jax.devices()[0]) + self.assertIsNone( + safetensors_layout._dedup_read_sharding(sharding, (8,)) + ) + + def test_no_replicating_axes_returns_none(self): + # Every mesh axis is used by the spec, so nothing is replicated and + # there is nothing to de-duplicate. + sharding = jax.sharding.NamedSharding( + jax.sharding.Mesh(np.array(jax.devices()), ('x',)), + jax.sharding.PartitionSpec('x'), + ) + self.assertIsNone( + safetensors_layout._dedup_read_sharding(sharding, (8,)) + ) + class SingleHostOneReadInvariantTest( unittest.IsolatedAsyncioTestCase, parameterized.TestCase