A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Core] feat: Add --safetensors-load-strategy flag for faster safetensors loading from Lustre (#24469)
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>
shengshiqi-google committed
Sep 11, 2025 at 06:10 UTC
41329a0ff95e916670e1a559e82d0d0d45d0883b
4 files changed
+31
-6
vllm/config/load.py
+9
@@ -51,6 +51,15 @@ class LoadConfig:
51
download_dir: Optional[str] = None
52
"""Directory to download and load the weights, default to the default
53
cache directory of Hugging Face."""
54
+ safetensors_load_strategy: Optional[str] = "lazy"
55
+ """Specifies the loading strategy for safetensors weights.
56
+ - "lazy" (default): Weights are memory-mapped from the file. This enables
57
+ on-demand loading and is highly efficient for models on local storage.
58
+ - "eager": The entire file is read into CPU memory upfront before loading.
59
+ This is recommended for models on network filesystems (e.g., Lustre, NFS)
60
+ as it avoids inefficient random reads, significantly speeding up model
61
+ initialization. However, it uses more CPU RAM.
62
+ """
63
model_loader_extra_config: Union[dict, TensorizerConfig] = field(
64
default_factory=dict)
65
"""Extra config for model loader. This will be passed to the model loader
vllm/engine/arg_utils.py
+5
@@ -289,6 +289,8 @@ class EngineArgs:
289
trust_remote_code: bool = ModelConfig.trust_remote_code
290
allowed_local_media_path: str = ModelConfig.allowed_local_media_path
291
download_dir: Optional[str] = LoadConfig.download_dir
292
+ safetensors_load_strategy: Optional[
293
+ str] = LoadConfig.safetensors_load_strategy
294
load_format: Union[str, LoadFormats] = LoadConfig.load_format
295
config_format: str = ModelConfig.config_format
296
dtype: ModelDType = ModelConfig.dtype
@@ -587,6 +589,8 @@ class EngineArgs:
589
load_group.add_argument("--load-format", **load_kwargs["load_format"])
590
load_group.add_argument("--download-dir",
591
**load_kwargs["download_dir"])
592
+ load_group.add_argument("--safetensors-load-strategy",
593
+ **load_kwargs["safetensors_load_strategy"])
594
load_group.add_argument("--model-loader-extra-config",
595
**load_kwargs["model_loader_extra_config"])
596
load_group.add_argument("--ignore-patterns",
@@ -1023,6 +1027,7 @@ class EngineArgs:
1027
return LoadConfig(
1028
load_format=self.load_format,
1029
download_dir=self.download_dir,
1030
+ safetensors_load_strategy=self.safetensors_load_strategy,
1031
device="cpu"
1032
if is_online_quantization(self.quantization) else None,
1033
model_loader_extra_config=self.model_loader_extra_config,
vllm/model_executor/model_loader/default_loader.py
+1
@@ -189,6 +189,7 @@ class DefaultModelLoader(BaseModelLoader):
189
weights_iterator = safetensors_weights_iterator(
190
hf_weights_files,
191
self.load_config.use_tqdm_on_load,
192
+ self.load_config.safetensors_load_strategy,
193
)
194
else:
195
if extra_config.get("enable_multithread_load"):
vllm/model_executor/model_loader/weight_utils.py
+16
-6
@@ -19,7 +19,7 @@ import huggingface_hub.constants
19
import numpy as np
20
import torch
21
from huggingface_hub import HfFileSystem, hf_hub_download, snapshot_download
22
-from safetensors.torch import load_file, safe_open, save_file
22
+from safetensors.torch import load, load_file, safe_open, save_file
23
from tqdm.auto import tqdm
24
25
from vllm import envs
@@ -519,18 +519,28 @@ def np_cache_weights_iterator(
519
def safetensors_weights_iterator(
520
hf_weights_files: list[str],
521
use_tqdm_on_load: bool,
522
+ safetensors_load_strategy: Optional[str] = "lazy",
523
) -> Generator[tuple[str, torch.Tensor], None, None]:
524
"""Iterate over the weights in the model safetensor files."""
525
+ loading_desc = "Loading safetensors checkpoint shards"
526
+ if safetensors_load_strategy == "eager":
527
+ loading_desc += " (eager)"
528
+
529
for st_file in tqdm(
530
hf_weights_files,
526
- desc="Loading safetensors checkpoint shards",
531
+ desc=loading_desc,
532
disable=not enable_tqdm(use_tqdm_on_load),
533
bar_format=_BAR_FORMAT,
534
):
530
- with safe_open(st_file, framework="pt") as f:
531
- for name in f.keys(): # noqa: SIM118
532
- param = f.get_tensor(name)
533
- yield name, param
535
+ if safetensors_load_strategy == "eager":
536
+ with open(st_file, "rb") as f:
537
+ state_dict = load(f.read())
538
+ yield from state_dict.items()
539
+ else:
540
+ with safe_open(st_file, framework="pt") as f:
541
+ for name in f.keys(): # noqa: SIM118
542
+ param = f.get_tensor(name)
543
+ yield name, param
544
545
546
def multi_thread_safetensors_weights_iterator(