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

[Perf] Use safetensors load_file in multithread loader (#18124)

Browse Source
This commit is contained in:
Mohammad Miadh Angkad
2026-02-02 23:21:13 -08:00
committed by GitHub
parent 7a9d9c79d1
commit 6f6b9c6e42
1 changed files with 2 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/model_loader/weight_utils.py
+2 -6
View File
@@ -840,12 +840,8 @@ def multi_thread_safetensors_weights_iterator(
def _load_file(st_file: str): def _load_file(st_file: str):
if disable_mmap: if disable_mmap:
with open(st_file, "rb") as f: with open(st_file, "rb") as f:
result = safetensors.torch.load(f.read()) return safetensors.torch.load(f.read())
else: return safetensors.torch.load_file(st_file, device="cpu")
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
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
futures = [executor.submit(_load_file, st_file) for st_file in hf_weights_files] futures = [executor.submit(_load_file, st_file) for st_file in hf_weights_files]