How do I load a safetensors checkpoint into sharded JAX arrays?

Give the loader an abstract state: a flat dict mapping each tensor name to a jax.ShapeDtypeStruct that carries the target jax.sharding.Sharding. Each process then resolves which shards its own devices own, maps those shards to byte ranges in the file, and reads exactly those ranges — the result is a sharded jax.Array with no all-gather and no full-file read.

Why is multihost checkpoint loading slow?

The naive approach has every host read the whole file, so a cluster of N hosts pulls N times the checkpoint size from storage. On object storage that is both egress cost and request-rate pressure. The sharding-driven loader reads only each host’s bytes; its slow case is strided inner-dimension sharding, which fragments into many tiny ranged reads — the loader coalesces those under a bounded over-read cap.

How does sharding decide which bytes a host reads?

A safetensors tensor is one contiguous row-major blob. A device’s shard is a slice of the tensor; that slice is contiguous in the file only along trailing dimensions. So a shard becomes a set of equally sized, equally spaced byte runs: exactly one run when it splits only the leading dimension (FSDP), and several strided runs when it splits an inner dimension (tensor parallelism).

What is the broadcast_replicated option in Orbax safetensors loading?

It removes the one case where sharding-driven reads still over-read: tensors replicated across processes, whose bytes every process would otherwise re-read. With the flag on, replicated tensors are read once under a de-duplicated sharding (each byte owned by one host, spread across the replicas), then an identity XLA computation broadcasts the shards to every replica over the accelerator interconnect — trading storage egress for interconnect traffic.

Does this bound host memory when loading a huge checkpoint?

Yes. Coalesced blocks are split at a fixed chunk size (128 MiB by default) and read concurrently, with total in-flight bytes capped by a shared budget (MemoryOptions.read_concurrent_bytes, default 2 GiB). Peak host memory beyond the destination shard buffers is bounded at that budget regardless of block, shard, or file size, and each file’s tensors are assembled and their buffers freed while other files are still reading.

Loading a Safetensors Checkpoint on a Multihost Cluster, Line by Line

A safetensors file is one flat byte blob per tensor, but a big model must land as sharded jax.Arrays across many hosts. Walk the real loader that maps each host’s shards to byte ranges — and reads exactly those.

Code walk · AI / ML. The source ↗

A free, interactive, animated visual explainer of Loading a Safetensors Checkpoint on a Multihost Cluster, Line by Line — built to be understood, not skimmed.

Questions

How do I load a safetensors checkpoint into sharded JAX arrays?: Give the loader an abstract state: a flat dict mapping each tensor name to a jax.ShapeDtypeStruct that carries the target jax.sharding.Sharding. Each process then resolves which shards its own devices own, maps those shards to byte ranges in the file, and reads exactly those ranges — the result is a sharded jax.Array with no all-gather and no full-file read.
Why is multihost checkpoint loading slow?: The naive approach has every host read the whole file, so a cluster of N hosts pulls N times the checkpoint size from storage. On object storage that is both egress cost and request-rate pressure. The sharding-driven loader reads only each host’s bytes; its slow case is strided inner-dimension sharding, which fragments into many tiny ranged reads — the loader coalesces those under a bounded over-read cap.
How does sharding decide which bytes a host reads?: A safetensors tensor is one contiguous row-major blob. A device’s shard is a slice of the tensor; that slice is contiguous in the file only along trailing dimensions. So a shard becomes a set of equally sized, equally spaced byte runs: exactly one run when it splits only the leading dimension (FSDP), and several strided runs when it splits an inner dimension (tensor parallelism).
What is the broadcast_replicated option in Orbax safetensors loading?: It removes the one case where sharding-driven reads still over-read: tensors replicated across processes, whose bytes every process would otherwise re-read. With the flag on, replicated tensors are read once under a de-duplicated sharding (each byte owned by one host, spread across the replicas), then an identity XLA computation broadcasts the shards to every replica over the accelerator interconnect — trading storage egress for interconnect traffic.
Does this bound host memory when loading a huge checkpoint?: Yes. Coalesced blocks are split at a fixed chunk size (128 MiB by default) and read concurrently, with total in-flight bytes capped by a shared budget (MemoryOptions.read_concurrent_bytes, default 2 GiB). Peak host memory beyond the destination shard buffers is bounded at that budget regardless of block, shard, or file size, and each file’s tensors are assembled and their buffers freed while other files are still reading.

Related explainers

Design a Distributed Message Queue