chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
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>
This commit is contained in:
Lianmin Zheng
2026-02-13 12:22:50 -08:00
committed by GitHub
parent 4c6afbeeaa
commit 008ea46af1
2 changed files with 68 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
python/sglang/srt/model_loader/loader.py
View File

@@ -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

111
python/sglang/srt/model_loader/weight_utils.py
View File

@@ -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 = (