From 008ea46af13e496240680932ad52c21b9e0589fb Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Fri, 13 Feb 2026 12:22:50 -0800 Subject: [PATCH] [Auto Sync] Update loader.py, weight_utils.py (20260213) (#18779) Co-authored-by: github-actions[bot] Co-authored-by: Xiuyu Li Co-authored-by: Cheng Wan <54331508+ch-wan@users.noreply.github.com> --- python/sglang/srt/model_loader/loader.py | 4 +- .../sglang/srt/model_loader/weight_utils.py | 111 +++++++++++------- 2 files changed, 68 insertions(+), 47 deletions(-) diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index a833682e5..da5c315af 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/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 diff --git a/python/sglang/srt/model_loader/weight_utils.py b/python/sglang/srt/model_loader/weight_utils.py index 17055e1c7..dac219058 100644 --- a/python/sglang/srt/model_loader/weight_utils.py +++ b/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. - - 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 - + """Iterate over the weights in the model safetensor files.""" 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(): - yield name, param + for name in sorted(result.keys()): + 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, - decryption_key: Optional[str] = None, - max_workers: int = 4, + max_workers: int, disable_mmap: bool = False, ) -> Generator[Tuple[str, torch.Tensor], None, None]: - """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 - + """Multi-Thread iterate over the weights in the model safetensor files.""" 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 = (