Prefetch only the byte spans each rank will read

[qgallouedec]

Stacked on #48227, which warms whole shards on every node and so carries the caveat "requires node RAM >= checkpoint size". It doesn't have to: every rank slices its own shard out of those files, so a node only needs the spans its ranks read.

Given the model's meta state dict, prefetch_checkpoint_shards reads the safetensors headers and the DTensor placements and warms only those ranges. Reading less costs seeks, so _SEEK_COST_BYTES prices a seek in bytes and whole files win whenever they should.

Paired cold runs, same nodes, differing only in whether the meta state dict is passed:

820B-A42B MoE, 1.5 TB, 8 nodes	bytes/rank	from_pretrained
whole shards	155-221 GiB	2455.7 s
rank spans	28.7 GiB	478.7 s

Node RAM now holds 229 GiB instead of 1.5 TB. Peak GPU memory 66.2 GB and median step 12.8 s in both arms.

The 12 MiB constant is measured, not chosen (right panel): reading a fixed 12.5% of a shard at varying span size, cache evicted per point. Below ~1.8 MiB spans, reading 12.5% is slower than reading all of it, which puts a seek at ~50 ms, or 12 MiB of sequential read. This is a cross-region mount; local NVMe is well below that, where the only effect is that whole-file warming is chosen more readily.

Checkpoints storing one tensor per expert (Qwen3-30B-A3B) make thousands of tiny spans, the cost test rejects them, and jobs becomes exactly whole[local_rank::local_world] — the pre-PR path, byte for byte.

Repro

CPU only, no download, seconds. Prints what one node warms for both layouts as the node count grows.

torchrun --nproc_per_node 4 repro_spans.py

import os, shutil, torch, torch.distributed as dist
from safetensors.torch import save_file
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor import Shard, distribute_tensor
from transformers.distributed.utils import _SEEK_COST_BYTES, _rank_byte_spans

EXPERTS, HIDDEN, LAYERS, DIRECTORY = 32, 2048, 2, "spans_repro_tmp"

def report(path, state_dict, rank, world, local_world, label):
    own, common = _rank_byte_spans([path], state_dict)
    jobs = own + common[rank % local_world :: local_world]
    needed = sum(end - start for _, start, end in jobs)
    cost = needed + len(jobs) * _SEEK_COST_BYTES
    whole = os.path.getsize(path) / local_world
    if rank == 0:
        print(f"  {label:11s} {world // local_world} node(s): warms {min(needed, whole) / 2**20:7.1f} MiB"
              f"  (spans {needed / 2**20:7.1f} MiB in {len(jobs):3d}, cost {cost / 2**20:8.1f} MiB"
              f"  vs whole files {whole / 2**20:7.1f} MiB) -> {'SPANS' if jobs and cost < whole else 'whole files'}", flush=True)

dist.init_process_group("gloo")
rank, world = dist.get_rank(), dist.get_world_size()
mesh = init_device_mesh("cpu", (world,))
packed, per_expert = f"{DIRECTORY}/packed.safetensors", f"{DIRECTORY}/per_expert.safetensors"
if rank == 0:
    os.makedirs(DIRECTORY, exist_ok=True)
    save_file({f"model.layers.{i}.mlp.experts.gate_up_proj":
               torch.zeros(EXPERTS, HIDDEN, HIDDEN, dtype=torch.bfloat16) for i in range(LAYERS)}, packed)
    save_file({f"model.layers.{i}.mlp.experts.{e}.gate_up_proj":
               torch.zeros(HIDDEN, HIDDEN, dtype=torch.bfloat16) for i in range(LAYERS) for e in range(EXPERTS)}, per_expert)
dist.barrier()

state_dict = {f"model.layers.{i}.mlp.experts.gate_up_proj":
              distribute_tensor(torch.zeros(EXPERTS, HIDDEN, HIDDEN, dtype=torch.bfloat16), mesh, [Shard(0)])
              for i in range(LAYERS)}
if rank == 0:
    print(f"\n{world} ranks, {os.path.getsize(packed) / 2**20:.0f} MiB checkpoint\n")
for local_world in [w for w in (world, world // 2, world // 4) if w >= 1]:
    report(packed, state_dict, rank, world, local_world, "packed")
    report(per_expert, state_dict, rank, world, local_world, "per-expert")
dist.barrier()
if rank == 0:
    shutil.rmtree(DIRECTORY, ignore_errors=True)
dist.destroy_process_group()

4 ranks, 512 MiB checkpoint

  packed      1 node(s): warms   128.0 MiB  (spans   128.0 MiB in   2, cost  152.0 MiB  vs whole files 128.0 MiB) -> whole files
  per-expert  1 node(s): warms   128.0 MiB  (spans   512.0 MiB in   1, cost  524.0 MiB  vs whole files 128.0 MiB) -> whole files
  packed      2 node(s): warms   128.0 MiB  (spans   128.0 MiB in   2, cost  152.0 MiB  vs whole files 256.0 MiB) -> SPANS
  per-expert  2 node(s): warms   256.0 MiB  (spans   512.0 MiB in   1, cost  524.0 MiB  vs whole files 256.0 MiB) -> whole files
  packed      4 node(s): warms   128.0 MiB  (spans   128.0 MiB in   2, cost  152.0 MiB  vs whole files 512.0 MiB) -> SPANS
  per-expert  4 node(s): warms   512.0 MiB  (spans   512.0 MiB in   1, cost  524.0 MiB  vs whole files 512.0 MiB) -> whole files

One node saves nothing, because whole-file warming already splits shards across a node's ranks. The saving is local_world / world, so it grows with the node count, and only for packed layouts.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[VI-Arthur]

let' s talk a bit about this together, this seems like a non typical case where you already need the checkoints to be saved in a certain format no?

[VI-Arthur]

This looks like a case of read then shard (GPU0 reads file0 layers 0-10, GPU1 file 1 layers 10-20, etc) -> you split like Pipeline, so you saturate coms, fewer reads no?

[VI-Arthur]

safetensors/safetensors#828

[github-actions]

---

CI recap

Dashboard: View test results in Grafana
Latest run: 34273963871:1
Result: success | Jobs: 16 | Tests: 186,596 | Failures: 0 | Duration: 15h 31m