Warm the page cache before sharded loading (opt-in via HF_SHARD_PREFETCH)

src/transformers/distributed/utils.py

@@ -59,6 +59,36 @@ def _get_torch_distributed_world_size() -> int:
     return torch.distributed.get_world_size()
 
 
+def prefetch_checkpoint_shards(checkpoint_files: list[str]) -> None:
+    """Warm the page cache for the checkpoint shards before the per-tensor loading pass, opt-in via
+    `HF_SHARD_PREFETCH=<read threads per rank>`.
+
+    The per-tensor read pattern of sharded loading reads a network filesystem at well under 1 GiB/s
+    while large sequential reads sustain many times that; warming the page cache first makes the
+    actual load run at memory speed. Local ranks split the shard list between them (every node needs
+    the full checkpoint cached, since every rank slices tensors from all shards).
+    """
+    prefetch_threads = int(os.environ.get("HF_SHARD_PREFETCH", "0"))
+    if not checkpoint_files or not prefetch_threads:
+        return
+    import time
+    from concurrent.futures import ThreadPoolExecutor
+
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    local_world = int(os.environ.get("LOCAL_WORLD_SIZE", "1"))
+
+    def _warm(path, bufsize=16 * 2**20):
+        with open(path, "rb", buffering=0) as f:
+            while f.read(bufsize):
+                pass
+
+    prefetch_start = time.time()
+    with ThreadPoolExecutor(max_workers=prefetch_threads) as pool:
+        list(pool.map(_warm, checkpoint_files[local_rank::local_world]))
+    _distributed_barrier()
+    logger.warning_once(f"Prefetched {len(checkpoint_files)} checkpoint shards in {time.time() - prefetch_start:.0f}s")
+
+
 def is_local_dist_rank_0() -> bool:
     return _is_torch_distributed_initialized() and int(os.environ.get("LOCAL_RANK", "-1")) == 0
 

src/transformers/modeling_utils.py

@@ -43,7 +43,7 @@
 from torch.distributions import constraints
 from torch.utils.checkpoint import checkpoint
 
-from transformers.distributed.utils import is_dtensor
+from transformers.distributed.utils import is_dtensor, prefetch_checkpoint_shards
 
 from . import initialization as init
 from .configuration_utils import PreTrainedConfig
@@ -4372,6 +4372,8 @@ def _load_pretrained_model(
         # 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
 
+        prefetch_checkpoint_shards(checkpoint_files)
+
         if logger.level >= logging.WARNING:
             verify_tp_plan(expected_keys, getattr(model, "_tp_plan", None))