[Auto Sync] Update loader.py, weight_utils.py (20260213) (#18779)

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Xiuyu Li <xiuyu@x.ai>
Co-authored-by: Cheng Wan <54331508+ch-wan@users.noreply.github.com>

python/sglang/srt/model_loader/loader.py

@@ -88,6 +88,7 @@ DEFAULT_GPU_MEMORY_FRACTION_FOR_CALIBRATION = (
 )
 from sglang.srt.environ import envs
 from sglang.srt.model_loader.weight_utils import (
+    buffered_multi_thread_safetensors_weights_iterator,
     download_safetensors_index_file_from_hf,
     download_weights_from_hf,
     fastsafetensors_weights_iterator,
@@ -99,7 +100,6 @@ from sglang.srt.model_loader.weight_utils import (
     initialize_dummy_weights,
     maybe_add_mtp_safetensors,
     multi_thread_pt_weights_iterator,
-    multi_thread_safetensors_weights_iterator,
     np_cache_weights_iterator,
     pt_weights_iterator,
     safetensors_weights_iterator,
@@ -508,7 +508,7 @@ class DefaultModelLoader(BaseModelLoader):
                     hf_weights_files,
                 )
             elif use_multithread:
-                weights_iterator = multi_thread_safetensors_weights_iterator(
+                weights_iterator = buffered_multi_thread_safetensors_weights_iterator(
                     hf_weights_files,
                     max_workers=extra_config.get(
                         "num_threads", self.DEFAULT_NUM_THREADS

python/sglang/srt/model_loader/weight_utils.py

@@ -1,10 +1,12 @@
 # Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/model_executor/model_loader/weight_utils.py
 
 """Utilities for downloading and initializing model weights."""
+import collections
 import concurrent.futures
 import fnmatch
 import glob
 import hashlib
+import itertools
 import json
 import logging
 import os
@@ -709,35 +711,11 @@ def np_cache_weights_iterator(
         yield name, torch.from_numpy(param)
 
 
-def decrypt(fn, key):
-    raise NotImplementedError()
-
-
-def safetensors_encrypted_weights_iterator(
-    hf_weights_files: List[str],
-    is_all_weights_sharded: bool = False,
-    decryption_key: Optional[str] = None,
-):
-    raise NotImplementedError()
-
-
 def safetensors_weights_iterator(
     hf_weights_files: List[str],
-    is_all_weights_sharded: bool = False,
-    decryption_key: Optional[str] = None,
     disable_mmap: bool = False,
 ) -> Generator[Tuple[str, torch.Tensor], None, None]:
-    """Iterate over the weights in the model safetensor files.
+    """Iterate over the weights in the model safetensor files."""
-
-    If is_all_weights_sharded is True, it uses more optimize read by reading an
-    entire file instead of reading each tensor one by one.
-    """
-    if decryption_key:
-        yield from safetensors_encrypted_weights_iterator(
-            hf_weights_files, is_all_weights_sharded, decryption_key
-        )
-        return
-
     enable_tqdm = (
         not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0
     )
@@ -751,8 +729,8 @@ def safetensors_weights_iterator(
         if disable_mmap:
             with open(st_file, "rb") as f:
                 result = safetensors.torch.load(f.read())
-                for name, param in result.items():
+                for name in sorted(result.keys()):
-                    yield name, param
+                    yield name, result[name]
         else:
             with safetensors.safe_open(st_file, framework="pt", device="cpu") as f:
                 for name in f.keys():
@@ -816,25 +794,10 @@ def fastsafetensors_weights_iterator(
 
 def multi_thread_safetensors_weights_iterator(
     hf_weights_files: List[str],
-    is_all_weights_sharded: bool = False,
+    max_workers: int,
-    decryption_key: Optional[str] = None,
-    max_workers: int = 4,
     disable_mmap: bool = False,
 ) -> Generator[Tuple[str, torch.Tensor], None, None]:
-    """Multi-Thread iterate over the weights in the model safetensor files.
+    """Multi-Thread iterate over the weights in the model safetensor files."""
-
-    If is_all_weights_sharded is True, it uses more optimize read by reading an
-    entire file instead of reading each tensor one by one.
-    """
-    if decryption_key:
-        logger.warning(
-            "Multi-Thread loading is not working for encrypted safetensor weights."
-        )
-        yield from safetensors_encrypted_weights_iterator(
-            hf_weights_files, is_all_weights_sharded, decryption_key
-        )
-        return
-
     enable_tqdm = (
         not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0
     )
@@ -865,6 +828,64 @@ def multi_thread_safetensors_weights_iterator(
                 yield name, param
 
 
+def buffered_multi_thread_safetensors_weights_iterator(
+    hf_weights_files: List[str],
+    max_workers: int,
+    disable_mmap: bool = False,
+) -> Generator[Tuple[str, torch.Tensor], None, None]:
+    """Multi-threaded safetensor loader with bounded memory via a sliding window.
+
+    At most (max_workers + 1) shard files are in-flight at any time:
+    max_workers loading concurrently + 1 prefetched and ready to yield.
+    Peak CPU RAM ≈ (max_workers + 2) × shard_file_size.
+    """
+    enable_tqdm = (
+        not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0
+    )
+
+    def _load_file(st_file: str):
+        if disable_mmap:
+            with open(st_file, "rb") as f:
+                result = safetensors.torch.load(f.read())
+        else:
+            with safetensors.safe_open(st_file, framework="pt", device="cpu") as f:
+                result = {k: f.get_tensor(k) for k in f.keys()}
+        return result
+
+    # Sliding window: max_workers loading + 1 prefetched.
+    buffer_size = max_workers + 1
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
+        file_iter = iter(hf_weights_files)
+        pending: collections.deque = collections.deque()
+
+        # Seed the buffer.
+        for st_file in itertools.islice(file_iter, buffer_size):
+            pending.append(executor.submit(_load_file, st_file))
+
+        with tqdm(
+            total=len(hf_weights_files),
+            desc="Multi-thread loading shards",
+            disable=not enable_tqdm,
+            bar_format=BAR_FORMAT,
+            position=tqdm._get_free_pos(),
+        ) as pbar:
+            while pending:
+                future = pending.popleft()
+                state_dict = future.result()
+                del future  # let GC reclaim the Future's internal result
+
+                # Replenish: submit the next file to keep the buffer full.
+                next_file = next(file_iter, None)
+                if next_file is not None:
+                    pending.append(executor.submit(_load_file, next_file))
+
+                for name in sorted(state_dict.keys()):
+                    yield name, state_dict[name]
+                del state_dict
+                pbar.update(1)
+
+
 def _load_pt_file(bin_file: str) -> dict:
     """Load a PyTorch checkpoint file, handling legacy tar format.
 
@@ -906,7 +927,7 @@ def pt_weights_iterator(
 
 def multi_thread_pt_weights_iterator(
     hf_weights_files: List[str],
-    max_workers: int = 4,
+    max_workers: int,
 ) -> Generator[Tuple[str, torch.Tensor], None, None]:
     """Multi-Thread iterate over the weights in the model bin/pt files."""
     enable_tqdm = (