Drive safetensors load by the target sharding with bounded concurrent ranged reads

[mridul-sahu]

Description

Rewrites the v1 SafetensorsLayout load path to be driven by the target sharding: each process resolves which shards of each tensor its own devices need (Sharding.devices_indices_map), maps those index domains to byte ranges in the file, and reads exactly those ranges. Resharding happens purely through which bytes each process reads — there is no cross-process communication and no XLA compilation. This mirrors how the v0 array deserializer drives reads from the target sharding, with hand-computed byte ranges standing in for TensorStore's chunk index.

Read planning. Byte runs from every tensor a process needs in one file are coalesced under a bounded over-read policy: SafetensorsOptions.max_over_read_ratio (default 2.0) caps block_size / needed_bytes, so inner-dim (e.g. row-parallel TP) shardings can't amplify per-host reads to N × file size, while a 4 KiB gap floor prevents strided reads from shattering into microscopic requests. Coalesced blocks are split at a fixed 128 MiB chunk size and all chunks — across all blocks and files — are issued as concurrent ranged reads, so large blocks fetch at parallel-stream throughput instead of one long single-stream read. Every chunk's destination-buffer copies are precomputed, so chunks may complete in any order.

Memory bound. One shared LimitInFlightBytes (reusing _src/serialization/limits.py, keyed off MemoryOptions.read_concurrent_bytes, 2 GiB fallback) throttles all in-flight chunks together: peak host memory beyond the destination shard buffers is bounded by the budget regardless of block, shard, or file size. Tensors are assembled into jax.Arrays per file as each file's reads complete, overlapping host-to-device transfer with the remaining reads and releasing host buffers progressively. File headers (and the fail-fast truncation check) are fetched concurrently across files.

Generality. Anything devices_indices_map can express now loads — inner-dim partitions, multi-axis partitions, N-D tensors — replacing the previous transient-array + collective-sum path, which required leading-dimension divisibility by the local device count and staged whole tensors on their owner hosts (the OOM risk tracked in its TODO).

Behavior changes.

• SafetensorsOptions.ignore_load_sharding is removed (its only users were its own tests; the converter it anticipated does not exist yet).
• The single-process path no longer rejects checkpoints containing non-finite values; the previous behavior was inconsistent between the single- and multi-host paths, and inf/NaN payloads (masks, sentinels) are legitimate.

Observability. Per-host reads are reported via jax.monitoring (/jax/orbax/read/safetensors/{bytes_read,num_reads,storage_reads}), and a one-shot warning fires when the over-read cap is what fragmented a load, pointing at max_over_read_ratio.

Tested with the single-process suite, a 4-process × 2-devices-per-process multiprocess suite covering the sharded read paths (including partial-replication, subset loads, and a peak-memory regression test), and the safetensors load benchmarks.

Type of change

• Bug fix (non-breaking change that fixes an issue)
• New feature (non-breaking change that adds functionality)
• Breaking change (alters existing behavior or a public API)
• Docs / tests / internal tooling (no user-facing behavior change)

Checklist

• I have read the contribution guidelines.
• Tests cover this change and pass locally.
• Public APIs and non-trivial functions are documented.
• If this change is user-facing, I have updated CHANGELOG.md under the [Unreleased] section (or the relevant checkpoint/ / export/ changelog).