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 15, 17:15 UTC
[Core] Integrate `fastsafetensors` loader for loading model weights (#10647)
Signed-off-by: Manish Sethi <Manish.sethi1@ibm.com>
Manish Sethi committed
Mar 24, 2025 at 11:08 UTC
761702fd19d15f472a7d57a3adbe0e9b5ae0bf98
11 files changed
+152
-9
docs/source/models/extensions/fastsafetensor.md
new
+5
@@ -0,0 +1,5 @@
1
+Loading Model weights with fastsafetensors
2
+===================================================================
3
+
4
+Using fastsafetensor library enables loading model weights to GPU memory by leveraging GPU direct storage. See https://github.com/foundation-model-stack/fastsafetensors for more details.
5
+For enabling this feature, set the environment variable ``USE_FASTSAFETENSOR`` to ``true``
docs/source/models/extensions/index.md
+1
@@ -5,4 +5,5 @@
5
6
runai_model_streamer
7
tensorizer
8
+fastsafetensor
9
:::
requirements/test.in
+1
@@ -41,3 +41,4 @@ tritonclient==2.51.0
41
numpy < 2.0.0
42
runai-model-streamer==0.11.0
43
runai-model-streamer-s3==0.11.0
44
+fastsafetensors>=0.1.10
requirements/test.txt
+12
-1
@@ -67,6 +67,7 @@ click==8.1.7
67
# jiwer
68
# nltk
69
# ray
70
+ # typer
71
colorama==0.4.6
72
# via
73
# awscli
@@ -122,6 +123,8 @@ fastparquet==2024.11.0
123
# via genai-perf
124
fastrlock==0.8.2
125
# via cupy-cuda12x
126
+fastsafetensors==0.1.10
127
+ # via -r requirements/test.in
128
filelock==3.16.1
129
# via
130
# datasets
@@ -505,7 +508,9 @@ requests==2.32.3
508
responses==0.25.3
509
# via genai-perf
510
rich==13.9.4
508
- # via genai-perf
511
+ # via
512
+ # genai-perf
513
+ # typer
514
rouge-score==0.1.2
515
# via lm-eval
516
rpds-py==0.20.1
@@ -550,6 +555,8 @@ setuptools==75.8.0
555
# via
556
# pytablewriter
557
# torch
558
+shellingham==1.5.4
559
+ # via typer
560
six==1.16.0
561
# via
562
# python-dateutil
@@ -600,6 +607,7 @@ torch==2.6.0
607
# accelerate
608
# bitsandbytes
609
# encodec
610
+ # fastsafetensors
611
# lm-eval
612
# peft
613
# runai-model-streamer
@@ -654,6 +662,8 @@ typepy==1.3.2
662
# dataproperty
663
# pytablewriter
664
# tabledata
665
+typer==0.15.2
666
+ # via fastsafetensors
667
typing-extensions==4.12.2
668
# via
669
# huggingface-hub
@@ -663,6 +673,7 @@ typing-extensions==4.12.2
673
# pydantic
674
# pydantic-core
675
# torch
676
+ # typer
677
tzdata==2024.2
678
# via pandas
679
urllib3==2.2.3
setup.py
+1
@@ -680,6 +680,7 @@ setup(
680
install_requires=get_requirements(),
681
extras_require={
682
"tensorizer": ["tensorizer>=2.9.0"],
683
+ "fastsafetensors": ["fastsafetensors >= 0.1.10"],
684
"runai": ["runai-model-streamer", "runai-model-streamer-s3", "boto3"],
685
"audio": ["librosa", "soundfile"], # Required for audio processing
686
"video": ["decord"] # Required for video processing
tests/fastsafetensors_loader/__init__.py
tests/fastsafetensors_loader/test_fastsafetensors_loader.py
new
+22
@@ -0,0 +1,22 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+
3
+from vllm import SamplingParams
4
+from vllm.config import LoadFormat
5
+
6
+test_model = "openai-community/gpt2"
7
+
8
+prompts = [
9
+ "Hello, my name is",
10
+ "The president of the United States is",
11
+ "The capital of France is",
12
+ "The future of AI is",
13
+]
14
+# Create a sampling params object.
15
+sampling_params = SamplingParams(temperature=0.8, top_p=0.95, seed=0)
16
+
17
+
18
+def test_model_loader_download_files(vllm_runner):
19
+ with vllm_runner(test_model,
20
+ load_format=LoadFormat.FASTSAFETENSORS) as llm:
21
+ deserialized_outputs = llm.generate(prompts, sampling_params)
22
+ assert deserialized_outputs
tests/fastsafetensors_loader/test_weight_utils.py
new
+46
@@ -0,0 +1,46 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+
3
+import glob
4
+import tempfile
5
+
6
+import huggingface_hub.constants
7
+import torch
8
+
9
+from vllm.model_executor.model_loader.weight_utils import (
10
+ download_weights_from_hf, fastsafetensors_weights_iterator,
11
+ safetensors_weights_iterator)
12
+
13
+
14
+def test_fastsafetensors_model_loader():
15
+ with tempfile.TemporaryDirectory() as tmpdir:
16
+ huggingface_hub.constants.HF_HUB_OFFLINE = False
17
+ download_weights_from_hf("openai-community/gpt2",
18
+ allow_patterns=["*.safetensors"],
19
+ cache_dir=tmpdir)
20
+ safetensors = glob.glob(f"{tmpdir}/**/*.safetensors", recursive=True)
21
+ assert len(safetensors) > 0
22
+
23
+ fastsafetensors_tensors = {}
24
+ hf_safetensors_tensors = {}
25
+
26
+ for name, tensor in fastsafetensors_weights_iterator(
27
+ safetensors, True):
28
+ fastsafetensors_tensors[name] = tensor
29
+
30
+ for name, tensor in safetensors_weights_iterator(safetensors, True):
31
+ hf_safetensors_tensors[name] = tensor
32
+
33
+ assert len(fastsafetensors_tensors) == len(hf_safetensors_tensors)
34
+
35
+ for name, fastsafetensors_tensor in fastsafetensors_tensors.items():
36
+ fastsafetensors_tensor = fastsafetensors_tensor.to('cpu')
37
+ assert fastsafetensors_tensor.dtype == hf_safetensors_tensors[
38
+ name].dtype
39
+ assert fastsafetensors_tensor.shape == hf_safetensors_tensors[
40
+ name].shape
41
+ assert torch.all(
42
+ fastsafetensors_tensor.eq(hf_safetensors_tensors[name]))
43
+
44
+
45
+if __name__ == "__main__":
46
+ test_fastsafetensors_model_loader()
vllm/config.py
+1
@@ -1277,6 +1277,7 @@ class LoadFormat(str, enum.Enum):
1277
BITSANDBYTES = "bitsandbytes"
1278
MISTRAL = "mistral"
1279
RUNAI_STREAMER = "runai_streamer"
1280
+ FASTSAFETENSORS = "fastsafetensors"
1281
1282
1283
@dataclass
vllm/model_executor/model_loader/loader.py
+16
-8
@@ -49,9 +49,10 @@ from vllm.model_executor.model_loader.utils import (ParamMapping,
49
set_default_torch_dtype)
50
from vllm.model_executor.model_loader.weight_utils import (
51
download_safetensors_index_file_from_hf, download_weights_from_hf,
52
- filter_duplicate_safetensors_files, filter_files_not_needed_for_inference,
53
- get_gguf_extra_tensor_names, get_lock, gguf_quant_weights_iterator,
54
- initialize_dummy_weights, np_cache_weights_iterator, pt_weights_iterator,
52
+ fastsafetensors_weights_iterator, filter_duplicate_safetensors_files,
53
+ filter_files_not_needed_for_inference, get_gguf_extra_tensor_names,
54
+ get_lock, gguf_quant_weights_iterator, initialize_dummy_weights,
55
+ np_cache_weights_iterator, pt_weights_iterator,
56
runai_safetensors_weights_iterator, safetensors_weights_iterator)
57
from vllm.model_executor.utils import set_weight_attrs
58
from vllm.platforms import current_platform
@@ -275,7 +276,8 @@ class DefaultModelLoader(BaseModelLoader):
276
# Some quantized models use .pt files for storing the weights.
277
if load_format == LoadFormat.AUTO:
278
allow_patterns = ["*.safetensors", "*.bin"]
278
- elif load_format == LoadFormat.SAFETENSORS:
279
+ elif (load_format == LoadFormat.SAFETENSORS
280
+ or load_format == LoadFormat.FASTSAFETENSORS):
281
use_safetensors = True
282
allow_patterns = ["*.safetensors"]
283
elif load_format == LoadFormat.MISTRAL:
@@ -357,10 +359,16 @@ class DefaultModelLoader(BaseModelLoader):
359
self.load_config.use_tqdm_on_load,
360
)
361
elif use_safetensors:
360
- weights_iterator = safetensors_weights_iterator(
361
- hf_weights_files,
362
- self.load_config.use_tqdm_on_load,
363
- )
362
+ if self.load_config.load_format == LoadFormat.FASTSAFETENSORS:
363
+ weights_iterator = fastsafetensors_weights_iterator(
364
+ hf_weights_files,
365
+ self.load_config.use_tqdm_on_load,
366
+ )
367
+ else:
368
+ weights_iterator = safetensors_weights_iterator(
369
+ hf_weights_files,
370
+ self.load_config.use_tqdm_on_load,
371
+ )
372
else:
373
weights_iterator = pt_weights_iterator(
374
hf_weights_files,
vllm/model_executor/model_loader/weight_utils.py
+47
@@ -38,6 +38,14 @@ except (ImportError, OSError):
38
SafetensorsStreamer = runai_model_streamer.placeholder_attr(
39
"SafetensorsStreamer")
40
41
+try:
42
+ from fastsafetensors import SafeTensorsFileLoader, SingleGroup
43
+except ImportError:
44
+ fastsafetensors = PlaceholderModule("fastsafetensors")
45
+ SafeTensorsFileLoader = fastsafetensors.placeholder_attr(
46
+ "SafeTensorsFileLoader")
47
+ SingleGroup = fastsafetensors.placeholder_attr("SingleGroup")
48
+
49
logger = init_logger(__name__)
50
51
# use system-level temp directory for file locks, so that multiple users
@@ -452,6 +460,45 @@ def runai_safetensors_weights_iterator(
460
yield from streamer.get_tensors()
461
462
463
+def fastsafetensors_weights_iterator(
464
+ hf_weights_files: List[str],
465
+ use_tqdm_on_load: bool,
466
+) -> Generator[Tuple[str, torch.Tensor], None, None]:
467
+ """Iterate over the weights in the model safetensor files
468
+ using fastsafetensor library."""
469
+ if torch.distributed.is_initialized():
470
+ pg = torch.distributed.group.WORLD
471
+ else:
472
+ pg = SingleGroup()
473
+
474
+ device = torch.device(f'cuda:{pg.rank()}')
475
+ weight_files_sub_lists = [
476
+ hf_weights_files[i:i + pg.size()]
477
+ for i in range(0, len(hf_weights_files), pg.size())
478
+ ]
479
+
480
+ for f_list in tqdm(
481
+ weight_files_sub_lists,
482
+ desc="Loading safetensors using Fastsafetensor loader",
483
+ disable=not enable_tqdm(use_tqdm_on_load),
484
+ bar_format=_BAR_FORMAT,
485
+ ):
486
+ loader = SafeTensorsFileLoader(pg, device)
487
+ rank_file_map = {i: [f] for i, f in enumerate(f_list)}
488
+ loader.add_filenames(rank_file_map)
489
+ try:
490
+ fb = loader.copy_files_to_device()
491
+ try:
492
+ keys = list(fb.key_to_rank_lidx.keys())
493
+ for k in keys:
494
+ t = fb.get_tensor(k)
495
+ yield k, t
496
+ finally:
497
+ fb.close()
498
+ finally:
499
+ loader.close()
500
+
501
+
502
def pt_weights_iterator(
503
hf_weights_files: List[str],
504
use_tqdm_on_load: bool,