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).

How does the loader pick its over-read ratio?

By default it picks one per file from the planned runs. It counts how many reads survive after merging only sub-page gaps: if that is at most 1024, the ratio is 1.0 and it never over-reads — a contiguous shard is already a handful of reads. Otherwise it uses that file’s span/needed, the exact ratio that collapses a fragmented strided plan into whole-span reads, bounding the request count at the cost of reading the gaps. You can override it with SafetensorsOptions.max_over_read_ratio, and a one-shot warning fires when the automatic ratio drove more than 1.5× over-read.

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).
How does the loader pick its over-read ratio?: By default it picks one per file from the planned runs. It counts how many reads survive after merging only sub-page gaps: if that is at most 1024, the ratio is 1.0 and it never over-reads — a contiguous shard is already a handful of reads. Otherwise it uses that file’s span/needed, the exact ratio that collapses a fragmented strided plan into whole-span reads, bounding the request count at the cost of reading the gaps. You can override it with SafetensorsOptions.max_over_read_ratio, and a one-shot warning fires when the automatic ratio drove more than 1.5× over-read.
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

Questions

Related explainers