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, 01:45 UTC
[Feat] Adds runai distributed streamer (#27230)
Signed-off-by: bbartels <benjamin@bartels.dev> Signed-off-by: Benjamin Bartels <benjamin@bartels.dev> Co-authored-by: omer-dayan <omdayan@nvidia.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Benjamin Bartels committed
Oct 30, 2025 at 04:09 UTC
17d055f527d2bd5d39a1352e5161ed82345466ac
9 files changed
+39
-11
docker/Dockerfile
+1
-1
@@ -495,7 +495,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
495
else \
496
BITSANDBYTES_VERSION="0.46.1"; \
497
fi; \
498
- uv pip install --system accelerate hf_transfer modelscope "bitsandbytes>=${BITSANDBYTES_VERSION}" 'timm>=1.0.17' 'runai-model-streamer[s3,gcs]>=0.14.0'
498
+ uv pip install --system accelerate hf_transfer modelscope "bitsandbytes>=${BITSANDBYTES_VERSION}" 'timm>=1.0.17' 'runai-model-streamer[s3,gcs]>=0.15.0'
499
500
ENV VLLM_USAGE_SOURCE production-docker-image
501
docs/models/extensions/runai_model_streamer.md
+9
@@ -45,6 +45,15 @@ vllm serve s3://core-llm/Llama-3-8b \
45
46
You can tune parameters using `--model-loader-extra-config`:
47
48
+You can tune `distributed` that controls whether distributed streaming should be used. This is currently only possible on CUDA and ROCM devices. This can significantly improve loading times from object storage or high-throughput network fileshares.
49
+You can read further about Distributed streaming [here](https://github.com/run-ai/runai-model-streamer/blob/master/docs/src/usage.md#distributed-streaming)
50
+
51
+```bash
52
+vllm serve /home/meta-llama/Llama-3.2-3B-Instruct \
53
+ --load-format runai_streamer \
54
+ --model-loader-extra-config '{"distributed":true}'
55
+```
56
+
57
You can tune `concurrency` that controls the level of concurrency and number of OS threads reading tensors from the file to the CPU buffer.
58
For reading from S3, it will be the number of client instances the host is opening to the S3 server.
59
requirements/nightly_torch_test.txt
+1
-1
@@ -42,6 +42,6 @@ tritonclient==2.51.0
42
43
numba == 0.61.2 # Required for N-gram speculative decoding
44
numpy
45
-runai-model-streamer[s3,gcs]==0.14.0
45
+runai-model-streamer[s3,gcs]==0.15.0
46
fastsafetensors>=0.1.10
47
pydantic>=2.12 # 2.11 leads to error on python 3.13
requirements/rocm.txt
+1
-1
@@ -12,6 +12,6 @@ tensorizer==2.10.1
12
packaging>=24.2
13
setuptools>=77.0.3,<80.0.0
14
setuptools-scm>=8
15
-runai-model-streamer[s3,gcs]==0.14.0
15
+runai-model-streamer[s3,gcs]==0.15.0
16
conch-triton-kernels==1.2.1
17
timm>=1.0.17
requirements/test.in
+1
-1
@@ -50,7 +50,7 @@ tritonclient==2.51.0
50
51
numba == 0.61.2 # Required for N-gram speculative decoding
52
numpy
53
-runai-model-streamer[s3,gcs]==0.14.0
53
+runai-model-streamer[s3,gcs]==0.15.0
54
fastsafetensors>=0.1.10
55
pydantic>=2.12 # 2.11 leads to error on python 3.13
56
decord==0.6.0
requirements/test.txt
+3
-3
@@ -965,11 +965,11 @@ rsa==4.9.1
965
# via google-auth
966
rtree==1.4.0
967
# via torchgeo
968
-runai-model-streamer==0.14.0
968
+runai-model-streamer==0.15.0
969
# via -r requirements/test.in
970
-runai-model-streamer-gcs==0.14.0
970
+runai-model-streamer-gcs==0.15.0
971
# via runai-model-streamer
972
-runai-model-streamer-s3==0.14.0
972
+runai-model-streamer-s3==0.15.0
973
# via runai-model-streamer
974
s3transfer==0.10.3
975
# via boto3
setup.py
+1
-1
@@ -712,7 +712,7 @@ setup(
712
"bench": ["pandas", "matplotlib", "seaborn", "datasets"],
713
"tensorizer": ["tensorizer==2.10.1"],
714
"fastsafetensors": ["fastsafetensors >= 0.1.10"],
715
- "runai": ["runai-model-streamer[s3,gcs] >= 0.14.0"],
715
+ "runai": ["runai-model-streamer[s3,gcs] >= 0.15.0"],
716
"audio": [
717
"librosa",
718
"soundfile",
vllm/model_executor/model_loader/runai_streamer_loader.py
+8
-2
@@ -27,9 +27,16 @@ class RunaiModelStreamerLoader(BaseModelLoader):
27
28
def __init__(self, load_config: LoadConfig):
29
super().__init__(load_config)
30
+
31
+ self._is_distributed = False
32
if load_config.model_loader_extra_config:
33
extra_config = load_config.model_loader_extra_config
34
35
+ if "distributed" in extra_config and isinstance(
36
+ extra_config.get("distributed"), bool
37
+ ):
38
+ self._is_distributed = extra_config.get("distributed")
39
+
40
if "concurrency" in extra_config and isinstance(
41
extra_config.get("concurrency"), int
42
):
@@ -92,8 +99,7 @@ class RunaiModelStreamerLoader(BaseModelLoader):
99
"""Get an iterator for the model weights based on the load format."""
100
hf_weights_files = self._prepare_weights(model_or_path, revision)
101
return runai_safetensors_weights_iterator(
95
- hf_weights_files,
96
- self.load_config.use_tqdm_on_load,
102
+ hf_weights_files, self.load_config.use_tqdm_on_load, self._is_distributed
103
)
104
105
def download_model(self, model_config: ModelConfig) -> None:
vllm/model_executor/model_loader/weight_utils.py
+14
-1
@@ -657,10 +657,22 @@ def multi_thread_safetensors_weights_iterator(
657
def runai_safetensors_weights_iterator(
658
hf_weights_files: list[str],
659
use_tqdm_on_load: bool,
660
+ is_distributed: bool = False,
661
) -> Generator[tuple[str, torch.Tensor], None, None]:
662
"""Iterate over the weights in the model safetensor files."""
663
with SafetensorsStreamer() as streamer:
663
- streamer.stream_files(hf_weights_files)
664
+ is_cuda_alike = current_platform.is_cuda_alike()
665
+ device = (
666
+ f"cuda:{current_platform.current_device()}"
667
+ if is_distributed and is_cuda_alike
668
+ else "cpu"
669
+ )
670
+
671
+ streamer.stream_files(
672
+ hf_weights_files,
673
+ device=device,
674
+ is_distributed=is_distributed,
675
+ )
676
total_tensors = sum(
677
len(tensors_meta)
678
for tensors_meta in streamer.files_to_tensors_metadata.values()
@@ -672,6 +684,7 @@ def runai_safetensors_weights_iterator(
684
desc="Loading safetensors using Runai Model Streamer",
685
bar_format=_BAR_FORMAT,
686
disable=not enable_tqdm(use_tqdm_on_load),
687
+ mininterval=2,
688
)
689
690
yield from tensor_iter