A high-throughput and memory-efficient inference and serving engine for LLMs
[Bugfix][Model] Pass revision by name in Run:ai and bitsandbytes index downloads (#45308)
Signed-off-by: Ting Sun <suntcrick@gmail.com>
Ting SUN committed
Jun 12, 2026 at 11:21 UTC
c1076839c9f14a51c0eb963ca8ae12c2de3c0f63
4 files changed
+59
-3
tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_loader.py
+25
@@ -1,11 +1,15 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import types
5
+from unittest.mock import patch
6
+
7
import pytest
8
9
from vllm import SamplingParams
10
from vllm.config.load import LoadConfig
11
from vllm.model_executor.model_loader import get_model_loader
12
+from vllm.model_executor.model_loader import runai_streamer_loader as rsl
13
14
load_format = "runai_streamer"
15
test_model = "openai-community/gpt2"
@@ -53,3 +57,24 @@ def test_runai_model_loader_download_files_gcs(
57
with vllm_runner(test_gcs_model, load_format=load_format) as llm:
58
deserialized_outputs = llm.generate(prompts, sampling_params)
59
assert deserialized_outputs
60
+
61
+
62
+def test_runai_passes_revision_by_name():
63
+ # revision must reach download_safetensors_index_file_from_hf as the
64
+ # ``revision`` keyword, not the positional ``subfolder`` slot.
65
+ fake_self = types.SimpleNamespace(
66
+ load_config=types.SimpleNamespace(download_dir="/cache", ignore_patterns=[])
67
+ )
68
+ with (
69
+ patch.object(rsl, "is_runai_obj_uri", return_value=False),
70
+ patch.object(rsl, "download_weights_from_hf", return_value="/folder"),
71
+ patch.object(
72
+ rsl, "list_safetensors", return_value=["/folder/model.safetensors"]
73
+ ),
74
+ patch.object(rsl, "download_safetensors_index_file_from_hf") as mock_idx,
75
+ ):
76
+ rsl.RunaiModelStreamerLoader._prepare_weights(fake_self, "org/model", "myrev")
77
+
78
+ mock_idx.assert_called_once()
79
+ assert mock_idx.call_args.kwargs.get("revision") == "myrev"
80
+ assert "myrev" not in mock_idx.call_args.args
tests/models/quantization/test_bitsandbytes.py
+28
@@ -5,12 +5,16 @@
5
Run `pytest tests/quantization/test_bitsandbytes.py`.
6
"""
7
8
+import types
9
+from unittest.mock import MagicMock, patch
10
+
11
import pytest
12
from packaging.version import Version
13
from transformers import BitsAndBytesConfig
14
from transformers import __version__ as TRANSFORMERS_VERSION
15
16
from tests.quantization.utils import is_quant_method_supported
17
+from vllm.model_executor.model_loader import bitsandbytes_loader as bnb
18
from vllm.platforms import current_platform
19
20
from ...utils import compare_two_settings, multi_gpu_test
@@ -300,3 +304,27 @@ def validate_generated_texts(
304
f"HF Output: '{hf_str}'\n"
305
f"vLLM Output: '{vllm_str}'"
306
)
307
+
308
+
309
+def test_bitsandbytes_passes_revision_by_name():
310
+ # revision must reach download_safetensors_index_file_from_hf as the
311
+ # ``revision`` keyword, not a positional slot.
312
+ fake_self = types.SimpleNamespace(
313
+ load_config=types.SimpleNamespace(download_dir="/cache"),
314
+ _get_weight_files=MagicMock(
315
+ return_value=("/folder", ["/folder/model.safetensors"], "*.safetensors")
316
+ ),
317
+ )
318
+ with (
319
+ patch.object(bnb, "download_safetensors_index_file_from_hf") as mock_idx,
320
+ patch.object(
321
+ bnb,
322
+ "filter_duplicate_safetensors_files",
323
+ return_value=["/folder/model.safetensors"],
324
+ ),
325
+ ):
326
+ bnb.BitsAndBytesModelLoader._prepare_weights(fake_self, "org/model", "myrev")
327
+
328
+ mock_idx.assert_called_once()
329
+ assert mock_idx.call_args.kwargs.get("revision") == "myrev"
330
+ assert "myrev" not in mock_idx.call_args.args
vllm/model_executor/model_loader/bitsandbytes_loader.py
+2
-2
@@ -140,8 +140,8 @@ class BitsAndBytesModelLoader(BaseModelLoader):
140
download_safetensors_index_file_from_hf(
141
model_name_or_path,
142
index_file,
143
- self.load_config.download_dir,
144
- revision,
143
+ cache_dir=self.load_config.download_dir,
144
+ revision=revision,
145
)
146
hf_weights_files = filter_duplicate_safetensors_files(
147
hf_weights_files, hf_folder, index_file
vllm/model_executor/model_loader/runai_streamer_loader.py
+4
-1
@@ -70,7 +70,10 @@ class RunaiModelStreamerLoader(BaseModelLoader):
70
71
if not is_local and not is_object_storage_path:
72
download_safetensors_index_file_from_hf(
73
- model_name_or_path, index_file, self.load_config.download_dir, revision
73
+ model_name_or_path,
74
+ index_file,
75
+ cache_dir=self.load_config.download_dir,
76
+ revision=revision,
77
)
78
79
if not hf_weights_files: