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, 01:00 UTC
Adds parallel model weight loading for runai_streamer (#21330)
Signed-off-by: bbartels <benjamin@bartels.dev> Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com>
Benjamin Bartels committed
Jul 22, 2025 at 16:15 UTC
b194557a6cfdd9eab777234c2ab3d90907e1c8f3
2 files changed
+16
-9
setup.py
+2
-1
@@ -659,7 +659,8 @@ setup(
659
"bench": ["pandas", "datasets"],
660
"tensorizer": ["tensorizer==2.10.1"],
661
"fastsafetensors": ["fastsafetensors >= 0.1.10"],
662
- "runai": ["runai-model-streamer", "runai-model-streamer-s3", "boto3"],
662
+ "runai":
663
+ ["runai-model-streamer >= 0.13.3", "runai-model-streamer-s3", "boto3"],
664
"audio": ["librosa", "soundfile",
665
"mistral_common[audio]"], # Required for audio processing
666
"video": [] # Kept for backwards compatibility
vllm/model_executor/model_loader/weight_utils.py
+14
-8
@@ -482,14 +482,20 @@ def runai_safetensors_weights_iterator(
482
) -> Generator[tuple[str, torch.Tensor], None, None]:
483
"""Iterate over the weights in the model safetensor files."""
484
with SafetensorsStreamer() as streamer:
485
- for st_file in tqdm(
486
- hf_weights_files,
487
- desc="Loading safetensors using Runai Model Streamer",
488
- disable=not enable_tqdm(use_tqdm_on_load),
489
- bar_format=_BAR_FORMAT,
490
- ):
491
- streamer.stream_file(st_file)
492
- yield from streamer.get_tensors()
485
+ streamer.stream_files(hf_weights_files)
486
+ total_tensors = sum(
487
+ len(tensors_meta)
488
+ for tensors_meta in streamer.files_to_tensors_metadata.values())
489
+
490
+ tensor_iter = tqdm(
491
+ streamer.get_tensors(),
492
+ total=total_tensors,
493
+ desc="Loading safetensors using Runai Model Streamer",
494
+ bar_format=_BAR_FORMAT,
495
+ disable=not enable_tqdm(use_tqdm_on_load),
496
+ )
497
+
498
+ yield from tensor_iter
499
500
501
def fastsafetensors_weights_iterator(