biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[Core] feat: Add --safetensors-load-strategy flag for faster safetensors loading from Lustre (#24469)

Browse Source 
Signed-off-by: Shiqi Sheng <shengshiqi@google.com>
Signed-off-by: shengshiqi-google <160179165+shengshiqi-google@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
shengshiqi-google
2025-09-11 06:10:01 +00:00
committed by GitHub
parent ee0bc5e1b4
commit 41329a0ff9
4 changed files with 31 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
vllm/model_executor/model_loader/default_loader.py
View File

@@ -189,6 +189,7 @@ class DefaultModelLoader(BaseModelLoader):
weights_iterator = safetensors_weights_iterator(
hf_weights_files,
self.load_config.use_tqdm_on_load,
self.load_config.safetensors_load_strategy,
)
else:
if extra_config.get("enable_multithread_load"):

22
vllm/model_executor/model_loader/weight_utils.py
View File

@@ -19,7 +19,7 @@ import huggingface_hub.constants
import numpy as np
import torch
from huggingface_hub import HfFileSystem, hf_hub_download, snapshot_download
from safetensors.torch import load_file, safe_open, save_file
from safetensors.torch import load, load_file, safe_open, save_file
from tqdm.auto import tqdm
from vllm import envs
@@ -519,18 +519,28 @@ def np_cache_weights_iterator(
def safetensors_weights_iterator(
hf_weights_files: list[str],
use_tqdm_on_load: bool,
safetensors_load_strategy: Optional[str] = "lazy",
) -> Generator[tuple[str, torch.Tensor], None, None]:
"""Iterate over the weights in the model safetensor files."""
loading_desc = "Loading safetensors checkpoint shards"
if safetensors_load_strategy == "eager":
loading_desc += " (eager)"
for st_file in tqdm(
hf_weights_files,
desc="Loading safetensors checkpoint shards",
desc=loading_desc,
disable=not enable_tqdm(use_tqdm_on_load),
bar_format=_BAR_FORMAT,
):
with safe_open(st_file, framework="pt") as f:
for name in f.keys(): # noqa: SIM118
param = f.get_tensor(name)
yield name, param
if safetensors_load_strategy == "eager":
with open(st_file, "rb") as f:
state_dict = load(f.read())
yield from state_dict.items()
else:
with safe_open(st_file, framework="pt") as f:
for name in f.keys(): # noqa: SIM118
param = f.get_tensor(name)
yield name, param
def multi_thread_safetensors_weights_iterator(