Warm the page cache before sharded loading (opt-in via HF_SHARD_PREFETCH) - #48227
Warm the page cache before sharded loading (opt-in via HF_SHARD_PREFETCH)#48227qgallouedec wants to merge 13 commits into
HF_SHARD_PREFETCH)#48227Conversation
…TCH) The loader's per-tensor read pattern pulls a network filesystem at well under 1 GiB/s while large sequential reads sustain many times that (measured: 0.26-0.7 GiB/s vs 8.5 GiB/s on Lustre). With HF_SHARD_PREFETCH=<threads>, the local ranks split the shard list and stream it into the page cache before loading; the load then runs at memory speed. Measured on GLM-4.5-Air (206 GiB, cold, 8 GPUs): 60 s baseline vs 23 s prefetch + 10 s load.
|
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. |
| # Model's definition arriving here is final (TP hooks added, quantized layers replaces) | ||
| expected_keys = list(model.state_dict().keys()) if expected_keys is None else expected_keys | ||
|
|
||
| if checkpoint_files and (prefetch_threads := int(os.environ.get("HF_SHARD_PREFETCH", "0"))): |
There was a problem hiding this comment.
very happy to have it, but in distributed/utils.py, isolated, imported.
ALSO we should default to something that makes sense if its always faster!
There was a problem hiding this comment.
@ArthurZucker happy to, but I'd rather base the default on more than our cluster.
What I measured: 2.6x cold on a shared FS, no measurable cost warm (+0.8 s on 93 s), but only 1.08x on node-local NVMe.
So it's a big win when the fs is remote and cold, neutral otherwise, and I haven't measured single-GPU or small-model loads.
Opened #48512 to track flipping it once we have numbers from other setups.
|
Hey @qgallouedec! Could you describe extremely precisely how you benchmarked this please? Because with our usuals cluster setups, what you will actually measure is usually simply whether or not weka already warmed up the files, not trying to prefetch them manually (the PR). If the files are warmed up by weka, everything is super fast. But the first run will be cold, and things will be slow. |
|
Note Edit: replaced the reproducer with a corrected version. @Cyrilvallez Fair, the first numbers didn't control for that. Redone with a protocol that removes the warm/cold confound: every run reads a checkpoint copy that has never been read on that node before, so neither the page cache nor the filesystem client can have it warm, and the two arms are interleaved so any drift in the storage backend hits both equally. (I also checked your weka-side concern directly, since a freshly written copy could plausibly still be hot in the backend: from a compute node, sequential read throughput on a fresh copy vs a file written 12 days ago and untouched since was 0.21 vs 0.18 GiB/s. So fresh writes get no meaningful server-side cache advantage here and both arms genuinely start cold.) Results on Qwen3-30B-A3B (57 GB, 8 ranks EP on one node),
The mechanism is only the read pattern: the per-tensor loading pass reads this filesystem at ~0.2 GiB/s single-stream, while a few large sequential readers sustain several times that. The prefetch doesn't beat a warm cache, it makes the cold path read at streaming speed instead of slice-at-a-time. Self-contained reproducer (makes its own copies, interleaves the arms, prints the table): bench_prefetch_repro.py"""Self-contained cold-cache benchmark for the shard prefetch (transformers #48227).
Every run reads a checkpoint copy that has never been read on this node, so neither the page
cache nor the filesystem client can have it warm; the two arms are interleaved so any drift in
the storage backend hits both equally.
# one node, 8 GPUs, ~6 x checkpoint size of free space
python bench_prefetch_repro.py --model Qwen/Qwen3-30B-A3B --scratch /shared/fs/bench --pairs 3
`--scratch` must live on the filesystem you actually load checkpoints from (the shared/network one).
Pointed at a node-local NVMe the whole question is moot: reads are already fast, and the measured
speedup drops to ~1.0x.
It prints one line per run and a summary table at the end.
"""
import argparse
import json
import os
import shutil
import statistics
import subprocess
import sys
import tempfile
import time
WORKER = r"""
import os, sys, time
os.environ["HF_SHARD_PREFETCH"] = sys.argv[2]
import torch
from transformers import AutoModelForCausalLM
from transformers.distributed import DistributedConfig
world = int(os.environ["WORLD_SIZE"])
t0 = time.time()
model = AutoModelForCausalLM.from_pretrained(
sys.argv[1],
dtype=torch.bfloat16,
distributed_config=DistributedConfig(tp_size=world, fsdp_size=1, enable_expert_parallel=True),
)
torch.distributed.barrier()
if int(os.environ["RANK"]) == 0:
print(f"LOAD_SECONDS {time.time() - t0:.1f}", flush=True)
"""
def run_once(worker_path, model_dir, prefetch, nproc, port):
cmd = [
"torchrun", "--nproc_per_node", str(nproc), "--master_port", str(port),
worker_path, model_dir, str(prefetch),
]
out = subprocess.run(cmd, capture_output=True, text=True)
for line in out.stdout.splitlines():
if line.startswith("LOAD_SECONDS"):
return float(line.split()[1])
sys.exit(f"run failed:\n{out.stdout[-2000:]}\n{out.stderr[-2000:]}")
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--model", default="Qwen/Qwen3-30B-A3B")
ap.add_argument(
"--scratch", required=True, help="on the shared filesystem under test; room for 2*pairs copies"
)
ap.add_argument("--pairs", type=int, default=3)
ap.add_argument("--prefetch-threads", type=int, default=4)
ap.add_argument("--nproc", type=int, default=8)
args = ap.parse_args()
from huggingface_hub import snapshot_download
source = snapshot_download(args.model)
os.makedirs(args.scratch, exist_ok=True)
# One never-read copy per run: the cold state is a property of the copy, not of an eviction trick.
copies = []
for i in range(2 * args.pairs):
dst = os.path.join(args.scratch, f"copy_{i}")
if not os.path.exists(dst):
print(f"copying {i + 1}/{2 * args.pairs} ...", flush=True)
shutil.copytree(source, dst, symlinks=False)
copies.append(dst)
with tempfile.NamedTemporaryFile("w", suffix=".py", delete=False) as f:
f.write(WORKER)
worker_path = f.name
results = {0: [], args.prefetch_threads: []}
for i, copy in enumerate(copies):
prefetch = 0 if i % 2 == 0 else args.prefetch_threads # interleaved arms
seconds = run_once(worker_path, copy, prefetch, args.nproc, 29500 + i)
results[prefetch].append(seconds)
print(f"### cold run {i + 1}: prefetch={prefetch} load={seconds:.1f}s", flush=True)
# Same copy twice more: the cache is warm now, so this is the cost of prefetching needlessly.
warm = copies[-1]
warm_off = run_once(worker_path, warm, 0, args.nproc, 29600)
warm_on = run_once(worker_path, warm, args.prefetch_threads, args.nproc, 29601)
os.unlink(worker_path)
off, on = results[0], results[args.prefetch_threads]
print("\n### RESULTS (load wall time, seconds)")
print(f"cold, no prefetch : {[round(x, 1) for x in off]} mean {statistics.mean(off):.1f}")
print(f"cold, prefetch={args.prefetch_threads} : {[round(x, 1) for x in on]} mean {statistics.mean(on):.1f}")
print(f"speedup : {statistics.mean(off) / statistics.mean(on):.2f}x")
print(f"warm, no prefetch : {warm_off:.1f}")
print(f"warm, prefetch={args.prefetch_threads} : {warm_on:.1f} (overhead {warm_on - warm_off:+.1f}s)")
print(json.dumps({"cold_off": off, "cold_on": on, "warm_off": warm_off, "warm_on": warm_on}))
if __name__ == "__main__":
main() |
CI recapDashboard: View test results in Grafana |
4 participants
Adds an opt-in page-cache prefetch to sharded model loading: with
HF_SHARD_PREFETCH=<threads>, the local ranks split the checkpoint's shard list between them and stream it into the page cache with large sequential reads before loading begins. The load itself then runs at memory speed.Why
The loader's per-tensor read pattern (mmap + slicing) reads a network filesystem far below what the hardware sustains. Measured on Lustre (H100 nodes, fadvise-evicted so genuinely cold):
So the entire gap is read scheduling, not conversion compute.
Measured
HF_SHARD_PREFETCH=4Multi-node totals are bounded by the filesystem's aggregate bandwidth shared across nodes (~10 GiB/s measured here): every node needs the full checkpoint cached, because with
distributed_configevery rank slices tensors from every shard. That is also why the env is opt-in: it requires node RAM ≥ checkpoint size.With the cache warm, the fixed 4-thread IO pool becomes the next bottleneck; raising it (already possible) halves the load again (Air: 39 s → 19 s across 8 ranks, skew < 3 s), and the usual raise-workers OOM hazard is gone because post-prefetch reads are cheap.
A note on the environments: the multi-node rows above come from remote nodes with ~15 ms of extra latency to the filesystem, which amplifies the loader's many-small-reads pattern; the single-node rows are low-latency nodes. Prefetch helps in both because large sequential reads are latency-tolerant.
Minimal repro
Runs on one node with 2+ GPUs and any sharded checkpoint (below: Qwen3-30B-A3B, 57 GiB, measured on 8×H100 with a Lustre-backed cache, 23 s cold baseline vs 13 s with prefetch):
Usage:
HF_SHARD_PREFETCH=4 torchrun --nproc_per_node 8 train.py # any from_pretrained with checkpoint filesFound while fine-tuning 100B–753B MoEs with FSDP2 × EP (#48204).
Part of the loading-performance work tracked in #48239.