@@ -8,6 +8,7 @@
import glob
import hashlib
import json
+import random
import os
import tempfile
import threading
@@ -880,6 +881,7 @@
use_tqdm_on_load: bool,
safetensors_load_strategy: str | None = None,
local_expert_ids: set[int] | None = None,
+ shuffle_safetensors_files: bool = False,
) -> Generator[tuple[str, torch.Tensor], None, None]:
"""Iterate over the weights in the model safetensor files.
@@ -892,6 +894,14 @@
loading_desc += " (eager)"
sorted_files = sorted(hf_weights_files, key=_natural_sort_key)
+
+ if shuffle_safetensors_files:
+ if torch.distributed.is_initialized():
+ rank = torch.distributed.get_rank()
+ else:
+ rank = 0
+ rng = random.Random(42 + rank)
+ rng.shuffle(sorted_files)
fs_type = _get_fs_type(sorted_files)
is_net_fs = fs_type in ("nfs", "nfs4", "lustre")
