@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 3d9a1d2de5
Fork 0

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

[V1] Support `LLM.apply_model` (#18465)

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

Cyrus Leung committed Sep 20, 2025 at 15:14 UTC 3d9a1d2de5091455bb2fbf6b21fc9188fd4612a4
17 files changed +194 -169
tests/conftest.py
+1 -11
@@ -987,17 +987,7 @@ class VllmRunner:
987 return [req_output.outputs.score for req_output in req_outputs]
988
989 def apply_model(self, func: Callable[[nn.Module], _R]) -> list[_R]:
990 - if hasattr(self.llm.llm_engine, "model_executor"):
991 - # This works either in V0 or in V1 with
992 - # VLLM_ENABLE_V1_MULTIPROCESSING=0
993 - executor = self.llm.llm_engine.model_executor
994 - return executor.apply_model(func)
995 -
996 - # This works in V1 with VLLM_ALLOW_INSECURE_SERIALIZATION=1
997 - def _apply_model(self):
998 - return func(self.get_model())
999 -
1000 - return self.llm.llm_engine.collective_rpc(_apply_model)
990 + return self.llm.apply_model(func)
991
992 def get_llm(self) -> LLM:
993 return self.llm
tests/kernels/moe/test_mxfp4_moe.py
+22 -15
@@ -1,21 +1,24 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -import importlib
4 import importlib.metadata
5 from dataclasses import dataclass
6 +from importlib.util import find_spec
7 from typing import Optional
8
9 import pytest
10 import torch
11 from packaging import version
12
13 +from vllm.model_executor.layers.quantization.quark.quark import ( # noqa: E501
14 + QuarkLinearMethod, QuarkW4A4MXFP4)
15 +from vllm.model_executor.layers.quantization.quark.quark_moe import ( # noqa: E501
16 + QuarkW4A4MXFp4MoEMethod)
17 from vllm.platforms import current_platform
18 from vllm.utils.flashinfer import has_flashinfer
19
16 -QUARK_MXFP4_AVAILABLE = importlib.util.find_spec(
17 - "quark") is not None and version.parse(
18 - importlib.metadata.version("amd-quark")) >= version.parse('0.8.99')
20 +QUARK_MXFP4_AVAILABLE = find_spec("quark") is not None and version.parse(
21 + importlib.metadata.version("amd-quark")) >= version.parse('0.8.99')
22
23 TRTLLM_GEN_MXFP4_AVAILABLE = current_platform.is_cuda(
24 ) and current_platform.is_device_capability(100)
@@ -39,6 +42,12 @@ class ModelCase:
42 tp: int
43
44
45 +@pytest.fixture(scope="function", autouse=True)
46 +def enable_pickle(monkeypatch):
47 + """`LLM.apply_model` requires pickling a function."""
48 + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
49 +
50 +
51 @pytest.mark.parametrize('model_case', [
52 ModelCase("fxmarty/qwen_1.5-moe-a2.7b-mxfp4", tp=1),
53 ModelCase("fxmarty/deepseek_r1_3_layers_mxfp4", tp=8),
@@ -55,21 +64,19 @@ def test_mxfp4_loading_and_execution_moe(vllm_runner, model_case: ModelCase):
64 tensor_parallel_size=model_case.tp,
65 load_format="dummy") as llm:
66
58 - # TODO: llm.apply_model(check_model) currently relies on V0 internals.
59 - # Re-enable this later.
60 - # def check_model(model):
61 - # layer = model.model.layers[0]
67 + def check_model(model):
68 + layer = model.model.layers[0]
69
63 - # qkv_proj = layer.self_attn.qkv_proj
70 + qkv_proj = layer.self_attn.qkv_proj
71
65 - # assert isinstance(qkv_proj.quant_method, QuarkLinearMethod)
66 - # assert isinstance(qkv_proj.scheme, QuarkW4A4MXFP4)
72 + assert isinstance(qkv_proj.quant_method, QuarkLinearMethod)
73 + assert isinstance(qkv_proj.scheme, QuarkW4A4MXFP4)
74
68 - # assert isinstance(layer.mlp.experts.quant_method,
69 - # QuarkW4A4MXFp4MoEMethod)
75 + assert isinstance(layer.mlp.experts.quant_method,
76 + QuarkW4A4MXFp4MoEMethod)
77
71 - # if model_case.model_id == "fxmarty/qwen_1.5-moe-a2.7b-mxfp4":
72 - # llm.apply_model(check_model)
78 + if model_case.model_id == "fxmarty/qwen_1.5-moe-a2.7b-mxfp4":
79 + llm.apply_model(check_model)
80
81 output = llm.generate_greedy("Today I am in the French Alps and",
82 max_tokens=20)
tests/models/multimodal/generation/test_qwen2_vl.py
+23 -23
@@ -10,6 +10,7 @@ from PIL import Image
10
11 from vllm.multimodal.image import rescale_image_size
12 from vllm.multimodal.video import rescale_video_size, sample_frames_from_video
13 +from vllm.utils import set_default_torch_num_threads
14
15 from ....conftest import (IMAGE_ASSETS, VIDEO_ASSETS, PromptImageInput,
16 PromptVideoInput, VllmRunner)
@@ -17,11 +18,9 @@ from ...utils import check_logprobs_close
18
19
20 @pytest.fixture(scope="function", autouse=True)
20 -def use_v0_only(monkeypatch):
21 - """
22 - V1 Test: batch_make_xxxxx_embeddings calls a V0 internal
23 - """
24 - monkeypatch.setenv('VLLM_USE_V1', '0')
21 +def enable_pickle(monkeypatch):
22 + """`LLM.apply_model` requires pickling a function."""
23 + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
24
25
26 models = ["Qwen/Qwen2-VL-2B-Instruct"]
@@ -126,9 +125,8 @@ def batch_make_image_embeddings(
125 image_grid_thw_on_device = image_grid_thw.to(visual.device,
126 dtype=torch.int64)
127 return visual(pixel_values_on_device,
129 - grid_thw=image_grid_thw_on_device)
128 + grid_thw=image_grid_thw_on_device).cpu()
129
131 - # V1 Test: this calls a V0 internal.
130 image_embeds = torch.concat(llm.apply_model(get_image_embeds))
131
132 # split into original batches
@@ -210,7 +208,7 @@ def batch_make_video_embeddings(
208 video_grid_thw_on_device = video_grid_thw.to(visual.device,
209 dtype=torch.int64)
210 return visual(pixel_values_on_device,
213 - grid_thw=video_grid_thw_on_device)
211 + grid_thw=video_grid_thw_on_device).cpu()
212
213 # V1 Test: this calls a V0 internal.
214 video_embeds = torch.concat(llm.apply_model(get_image_embeds))
@@ -266,19 +264,22 @@ def run_embedding_input_test(
264 processor = AutoProcessor.from_pretrained(model)
265
266 # max_model_len should be greater than image_feature_size
269 - with vllm_runner(model,
270 - runner="generate",
271 - max_model_len=4000,
272 - max_num_seqs=3,
273 - dtype=dtype,
274 - limit_mm_per_prompt={
275 - "image": mm_limit,
276 - "video": mm_limit
277 - },
278 - tensor_parallel_size=tensor_parallel_size,
279 - distributed_executor_backend=distributed_executor_backend
280 - ) as vllm_model:
267 + with set_default_torch_num_threads(1):
268 + vllm_model = vllm_runner(
269 + model,
270 + runner="generate",
271 + max_model_len=4000,
272 + max_num_seqs=3,
273 + dtype=dtype,
274 + limit_mm_per_prompt={
275 + "image": mm_limit,
276 + "video": mm_limit
277 + },
278 + tensor_parallel_size=tensor_parallel_size,
279 + distributed_executor_backend=distributed_executor_backend,
280 + )
281
282 + with vllm_model:
283 outputs_per_case_for_original_input = [
284 vllm_model.generate_greedy_logprobs(prompts,
285 max_tokens,
@@ -329,9 +330,8 @@ def run_embedding_input_test(
330 @pytest.mark.parametrize("max_tokens", [128])
331 @pytest.mark.parametrize("num_logprobs", [10])
332 def test_qwen2_vl_image_embeddings_input(vllm_runner, image_assets, model,
332 - size_factors, dtype: str,
333 - max_tokens: int,
334 - num_logprobs: int) -> None:
333 + size_factors, dtype, max_tokens,
334 + num_logprobs, monkeypatch) -> None:
335 images = [asset.pil_image for asset in image_assets]
336
337 inputs_per_case: list[tuple[
tests/models/quantization/test_awq.py
+1 -1
@@ -112,7 +112,7 @@ def test_awq_models(vllm_runner, image_assets, source_model, quant_model,
112 monkeypatch) -> None:
113
114 # Test V1: this test hangs during setup on single-scale input.
115 - # TODO: fixure out why and re-enable this on V1.
115 + # TODO: figure out why and re-enable this on V1.
116 monkeypatch.setenv("VLLM_USE_V1", "0")
117 run_awq_test(
118 vllm_runner,
tests/quantization/test_compressed_tensors.py
+8 -10
@@ -43,12 +43,9 @@ ROCM_TRITON_SCALED_MM_SUPPORTED_INT8_MODEL = [
43
44
45 @pytest.fixture(scope="function", autouse=True)
46 -def use_v0_only(monkeypatch):
47 - """
48 - This module relies on V0 internals, so set VLLM_USE_V1=0.
49 - """
50 - if not current_platform.is_cpu():
51 - monkeypatch.setenv('VLLM_USE_V1', '0')
46 +def enable_pickle(monkeypatch):
47 + """`LLM.apply_model` requires pickling a function."""
48 + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
49
50
51 @pytest.mark.parametrize(
@@ -176,10 +173,11 @@ def test_compressed_tensors_w8a8_logprobs(
173
174 dtype = "bfloat16"
175
179 - # skip language translation prompt for the static per tensor asym model
180 - if (model_path ==
181 - "nm-testing/Meta-Llama-3-8B-Instruct-W8A8-Static-Per-Tensor-Asym"
182 - ): # noqa: E501
176 + # skip language translation prompt for the static per tensor models
177 + if model_path in (
178 + "nm-testing/Meta-Llama-3-8B-Instruct-W8A8-Static-Per-Tensor-Sym",
179 + "nm-testing/Meta-Llama-3-8B-Instruct-W8A8-Static-Per-Tensor-Asym",
180 + ):
181 example_prompts = example_prompts[0:-1]
182
183 with hf_runner(model_path, dtype=dtype) as hf_model:
tests/quantization/test_fp8.py
+4 -4
@@ -60,8 +60,8 @@ def test_kv_cache_model_load_and_run(vllm_runner, model_id: str,
60 if use_rocm_aiter:
61 monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
62
63 - # vllm_runner.apply_model() relies on V0 internals.
64 - monkeypatch.setenv("VLLM_USE_V1", "0")
63 + # `LLM.apply_model` requires pickling a function.
64 + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
65 with vllm_runner(model_id, kv_cache_dtype="fp8") as llm:
66
67 def check_model(model):
@@ -104,8 +104,8 @@ def test_load_fp16_model(vllm_runner, kv_cache_dtype: str, force_marlin: bool,
104 if use_rocm_aiter:
105 monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
106
107 - # vllm_runner.apply_model() relies on V0 internals.
108 - monkeypatch.setenv("VLLM_USE_V1", "0")
107 + # `LLM.apply_model` requires pickling a function.
108 + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
109
110 if force_marlin:
111 monkeypatch.setenv("VLLM_TEST_FORCE_FP8_MARLIN", "1")
tests/quantization/test_gptq_dynamic.py
+38 -33
@@ -31,41 +31,46 @@ MODEL_QUANT = [
31 @pytest.mark.parametrize("model_id, use_marlin_kernel", MODEL_QUANT)
32 def test_gptq_with_dynamic(vllm_runner, model_id: str, use_marlin_kernel: bool,
33 monkeypatch):
34 - # vllm_runner.apply_model() relies on V0 internals.
35 - monkeypatch.setenv("VLLM_USE_V1", "0")
36 -
37 - vllm_model = vllm_runner(model_id, dtype=torch.float16, max_model_len=2048)
34 + # `LLM.apply_model` requires pickling a function.
35 + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
36
37 linear_method_cls = GPTQMarlinLinearMethod if use_marlin_kernel else (
38 GPTQLinearMethod)
39
42 - for name, submodule in (vllm_model.llm.llm_engine.model_executor.
43 - driver_worker.model_runner.model.named_modules()):
44 - if name == "lm_head":
45 - assert isinstance(submodule.quant_method, linear_method_cls)
46 - elif name == 'model.layers.0.self_attn.qkv_proj':
47 - # The first layer is quantized using bits=4, group_size=128
48 - # desc_act=True
49 - assert isinstance(submodule.quant_method, linear_method_cls)
50 - config = submodule.quant_method.quant_config
51 - assert config.weight_bits == 4
52 - assert config.group_size == 128
53 - assert config.desc_act
54 - elif name == 'model.layers.1.self_attn.qkv_proj':
55 - # The second layer is quantized using bits=8, group_size=32
56 - # desc_act=False
57 - assert isinstance(submodule.quant_method, linear_method_cls)
58 - config = submodule.quant_method.quant_config
59 - assert get_dynamic_override(config, layer_name=name,
60 - key="bits") == 8
61 - assert get_dynamic_override(config,
62 - layer_name=name,
63 - key="group_size") == 32
64 - assert not get_dynamic_override(
65 - config, layer_name=name, key="desc_act")
66 - elif (name == 'model.layers.2.self_attn.qkv_proj'
67 - or name == 'model.layers.2.mlp.gate_up_proj'):
68 - # All other layers (layer index >= 2) are not quantized
69 - assert isinstance(submodule.quant_method, UnquantizedLinearMethod)
40 + with vllm_runner(model_id, dtype=torch.float16, max_model_len=2048) as llm:
41 +
42 + def check_model(model):
43 + for name, submodule in model.named_modules():
44 + if name == "lm_head":
45 + assert isinstance(submodule.quant_method,
46 + linear_method_cls)
47 + elif name == 'model.layers.0.self_attn.qkv_proj':
48 + # The first layer is quantized using bits=4, group_size=128
49 + # desc_act=True
50 + assert isinstance(submodule.quant_method,
51 + linear_method_cls)
52 + config = submodule.quant_method.quant_config
53 + assert config.weight_bits == 4
54 + assert config.group_size == 128
55 + assert config.desc_act
56 + elif name == 'model.layers.1.self_attn.qkv_proj':
57 + # The second layer is quantized using bits=8, group_size=32
58 + # desc_act=False
59 + assert isinstance(submodule.quant_method,
60 + linear_method_cls)
61 + config = submodule.quant_method.quant_config
62 + assert get_dynamic_override(config,
63 + layer_name=name,
64 + key="bits") == 8
65 + assert get_dynamic_override(config,
66 + layer_name=name,
67 + key="group_size") == 32
68 + assert not get_dynamic_override(
69 + config, layer_name=name, key="desc_act")
70 + elif (name == 'model.layers.2.self_attn.qkv_proj'
71 + or name == 'model.layers.2.mlp.gate_up_proj'):
72 + # All other layers (layer index >= 2) are not quantized
73 + assert isinstance(submodule.quant_method,
74 + UnquantizedLinearMethod)
75
71 - del vllm_model
76 + llm.apply_model(check_model)
tests/quantization/test_lm_head.py
+2 -2
@@ -29,8 +29,8 @@ def test_lm_head(
29 lm_head_quantized: bool,
30 monkeypatch,
31 ) -> None:
32 - # vllm_runner.apply_model() relies on V0 internals.
33 - monkeypatch.setenv("VLLM_USE_V1", "0")
32 + # `LLM.apply_model` requires pickling a function.
33 + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
34 with vllm_runner(model_id, dtype=torch.float16,
35 max_model_len=2048) as vllm_model:
36
tests/quantization/test_modelopt.py
+3 -7
@@ -11,16 +11,12 @@ import pytest
11 import torch
12
13 from tests.quantization.utils import is_quant_method_supported
14 -from vllm.platforms import current_platform
14
15
16 @pytest.fixture(scope="function", autouse=True)
18 -def use_v0_only(monkeypatch):
19 - """
20 - This module relies on V0 internals, so set VLLM_USE_V1=0.
21 - """
22 - if not current_platform.is_cpu():
23 - monkeypatch.setenv('VLLM_USE_V1', '0')
17 +def enable_pickle(monkeypatch):
18 + """`LLM.apply_model` requires pickling a function."""
19 + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
20
21
22 @pytest.mark.skipif(not is_quant_method_supported("modelopt"),
tests/quantization/test_ptpc_fp8.py
+28 -19
@@ -13,6 +13,16 @@ from vllm.model_executor.layers.quantization.ptpc_fp8 import (
13 PTPCFp8LinearMethod)
14 from vllm.platforms import current_platform
15
16 +UNSUPPORTED_STR = (
17 + "Currently torch._scaled_mm (hipBLASLt) rowwise gemm only "
18 + "support output dtype of bfloat16. torch.float16 is specified.")
19 +
20 +
21 +@pytest.fixture(scope="function", autouse=True)
22 +def enable_pickle(monkeypatch):
23 + """`LLM.apply_model` requires pickling a function."""
24 + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
25 +
26
27 @pytest.mark.skipif(not is_quant_method_supported("ptpc_fp8"),
28 reason="PTPC FP8 is not supported on this GPU type.")
@@ -21,14 +31,22 @@ from vllm.platforms import current_platform
31 @pytest.mark.parametrize("dtype", ["auto", "bfloat16", "float16"])
32 @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8", "fp8_e4m3"])
33 def test_ptpc_fp8_rocm(vllm_runner, dtype: str, kv_cache_dtype: str) -> None:
24 -
34 try:
26 - with vllm_runner("facebook/opt-125m",
27 - dtype=dtype,
28 - quantization="ptpc_fp8",
29 - kv_cache_dtype=kv_cache_dtype) as llm:
35 + llm = vllm_runner("facebook/opt-125m",
36 + dtype=dtype,
37 + quantization="ptpc_fp8",
38 + kv_cache_dtype=kv_cache_dtype)
39 + except AssertionError as e:
40 + if str(e) == UNSUPPORTED_STR:
41 + # If the error message matches, the test passes
42 + return
43 + else:
44 + # If the error message does not match, re-raise the exception
45 + raise
46 +
47 + with llm:
48
31 - model = llm.model.llm_engine.model_executor.driver_worker.model_runner.model # noqa: E501
49 + def check_model(model):
50 fc1 = model.model.decoder.layers[0].fc1
51 assert isinstance(fc1.quant_method, PTPCFp8LinearMethod)
52 if kv_cache_dtype == "ptpc_fp8":
@@ -40,17 +58,8 @@ def test_ptpc_fp8_rocm(vllm_runner, dtype: str, kv_cache_dtype: str) -> None:
58 if current_platform.has_device_capability(94):
59 # For GPUs with hardware support, we keep weights in fp8
60 assert fc1.weight.dtype == torch.float8_e4m3fnuz
43 - else:
44 - pytest.skip()
61
46 - output = llm.generate_greedy("Hello my name is", max_tokens=20)
47 - assert output
48 - except AssertionError as e:
49 - if str(
50 - e
51 - ) == "Currently torch._scaled_mm (hipBLASLt) rowwise gemm only support output dtype of bfloat16. torch.float16 is specified.": # noqa: E501
52 - # If the error message matches, the test passes
53 - pass
54 - else:
55 - # If the error message does not match, re-raise the exception
56 - raise
62 + llm.apply_model(check_model)
63 +
64 + output = llm.generate_greedy("Hello my name is", max_tokens=20)
65 + assert output
tests/quantization/test_quark.py
+11 -15
@@ -7,10 +7,10 @@ Run `pytest tests/quantization/test_quark.py`.
7 See also `tests/kernels/moe/test_mxfp4_moe.py`.
8 """
9
10 -import importlib
10 import importlib.metadata
11 import os
12 from dataclasses import dataclass
13 +from importlib.util import find_spec
14
15 import huggingface_hub
16 import lm_eval
@@ -24,9 +24,8 @@ from vllm.platforms import current_platform
24
25 from .reference_mxfp4 import dq_mxfp4_torch, qdq_mxfp4_torch
26
27 -QUARK_MXFP4_AVAILABLE = importlib.util.find_spec(
28 - "quark") is not None and version.parse(
29 - importlib.metadata.version("amd-quark")) >= version.parse('0.8.99')
27 +QUARK_MXFP4_AVAILABLE = find_spec("quark") is not None and version.parse(
28 + importlib.metadata.version("amd-quark")) >= version.parse('0.8.99')
29
30 if QUARK_MXFP4_AVAILABLE:
31 from quark.torch.export.nn.modules.realquantizer import (
@@ -43,11 +42,9 @@ except huggingface_hub.errors.RepositoryNotFoundError:
42
43
44 @pytest.fixture(scope="function", autouse=True)
46 -def use_v0_only(monkeypatch):
47 - """
48 - This module relies on V0 internals, so set VLLM_USE_V1=0.
49 - """
50 - monkeypatch.setenv('VLLM_USE_V1', '0')
45 +def enable_pickle(monkeypatch):
46 + """`LLM.apply_model` requires pickling a function."""
47 + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
48
49
50 @pytest.mark.parametrize('kv_cache_dtype', ['auto', 'fp8'])
@@ -132,13 +129,12 @@ def test_quark_fp8_parity(vllm_runner):
129 }
130 with (vllm_runner(quark_model_id, **llm_kwargs) as
131 quark_handle, vllm_runner(fp8_model_id, **llm_kwargs) as fp8_handle):
135 - quark_model = (quark_handle.llm.llm_engine.model_executor.
136 - driver_worker.model_runner.model)
137 - quark_state_dict = quark_model.state_dict()
132
139 - fp8_model = (fp8_handle.llm.llm_engine.model_executor.driver_worker.
140 - model_runner.model)
141 - fp8_state_dict = fp8_model.state_dict()
133 + def get_state_dict(model):
134 + return {k: v.cpu() for k, v in model.state_dict().items()}
135 +
136 + quark_state_dict, = quark_handle.apply_model(get_state_dict)
137 + fp8_state_dict, = fp8_handle.apply_model(get_state_dict)
138
139 assert fp8_state_dict.keys() == quark_state_dict.keys()
140
tests/quantization/test_register_quantization_config.py
+10 -7
@@ -105,18 +105,21 @@ def test_register_quantization_config():
105 ])
106 def test_custom_quant(vllm_runner, model, monkeypatch):
107 """Test infer with the custom quantization method."""
108 - # vllm_runner.apply_model() relies on V0 internals.
109 - monkeypatch.setenv("VLLM_USE_V1", "0")
108 + # `LLM.apply_model` requires pickling a function.
109 + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
110 +
111 with vllm_runner(model_name=model,
112 quantization="custom_quant",
113 enforce_eager=True) as llm:
114
114 - model = llm.llm.llm_engine.model_executor.driver_worker.model_runner.model # noqa: E501
115 - layer = model.model.layers[0]
116 - qkv_proj = layer.self_attn.qkv_proj
115 + def check_model(model):
116 + layer = model.model.layers[0]
117 + qkv_proj = layer.self_attn.qkv_proj
118 +
119 + # Check the quantization method is FakeQuantLinearMethod
120 + assert isinstance(qkv_proj.quant_method, FakeQuantLinearMethod)
121
118 - # Check the quantization method is FakeQuantLinearMethod
119 - assert isinstance(qkv_proj.quant_method, FakeQuantLinearMethod)
122 + llm.apply_model(check_model)
123
124 output = llm.generate_greedy("Hello my name is", max_tokens=20)
125 assert output
vllm/engine/llm_engine.py
+6 -1
@@ -13,6 +13,7 @@ from typing import Sequence as GenericSequence
13 from typing import Set, Type, Union, cast
14
15 import torch
16 +import torch.nn as nn
17 from typing_extensions import TypeVar
18
19 import vllm.envs as envs
@@ -55,6 +56,7 @@ from vllm.usage.usage_lib import (UsageContext, is_usage_stats_enabled,
56 from vllm.utils import Counter, Device, resolve_obj_by_qualname, weak_bind
57 from vllm.version import __version__ as VLLM_VERSION
58 from vllm.worker.model_runner_base import InputProcessingError
59 +from vllm.worker.worker_base import WorkerBase
60
61 logger = init_logger(__name__)
62 _LOCAL_LOGGING_INTERVAL_SEC = 5
@@ -1817,13 +1819,16 @@ class LLMEngine:
1819 return sampling_params
1820
1821 def collective_rpc(self,
1820 - method: Union[str, Callable[..., _R]],
1822 + method: Union[str, Callable[[WorkerBase], _R]],
1823 timeout: Optional[float] = None,
1824 args: tuple = (),
1825 kwargs: Optional[dict[str, Any]] = None) -> list[_R]:
1826 return self.model_executor.collective_rpc(method, timeout, args,
1827 kwargs)
1828
1829 + def apply_model(self, func: Callable[[nn.Module], _R]) -> list[_R]:
1830 + return self.collective_rpc("apply_model", args=(func, ))
1831 +
1832
1833 if envs.is_set("VLLM_USE_V1") and envs.VLLM_USE_V1:
1834 from vllm.v1.engine.llm_engine import LLMEngine as V1LLMEngine
vllm/entrypoints/llm.py
+7 -2
@@ -522,9 +522,14 @@ class LLM:
522 """
523 Run a function directly on the model inside each worker,
524 returning the result for each of them.
525 +
526 + !!! warning
527 + To reduce the overhead of data transfer, avoid returning large
528 + arrays or tensors from this method. If you must return them,
529 + make sure you move them to CPU first to avoid taking up additional
530 + VRAM!
531 """
526 - executor = self.llm_engine.model_executor
527 - return executor.apply_model(func)
532 + return self.llm_engine.apply_model(func)
533
534 def _get_beam_search_lora_requests(
535 self,
vllm/executor/executor_base.py
+16 -17
@@ -5,11 +5,10 @@ import asyncio
5 import time
6 from abc import ABC, abstractmethod
7 from functools import cached_property
8 -from typing import (Any, Awaitable, Callable, Dict, List, Optional, Set, Tuple,
9 - Union)
8 +from typing import Any, Awaitable, Callable, List, Optional, Set, Union
9
10 import torch.nn as nn
12 -from typing_extensions import TypeVar
11 +from typing_extensions import TypeVar, deprecated
12
13 import vllm.platforms
14 from vllm.config import VllmConfig
@@ -63,10 +62,10 @@ class ExecutorBase(ABC):
62
63 @abstractmethod
64 def collective_rpc(self,
66 - method: Union[str, Callable[..., _R]],
65 + method: Union[str, Callable[[WorkerBase], _R]],
66 timeout: Optional[float] = None,
68 - args: Tuple = (),
69 - kwargs: Optional[Dict[str, Any]] = None) -> List[_R]:
67 + args: tuple = (),
68 + kwargs: Optional[dict[str, Any]] = None) -> list[_R]:
69 """
70 Execute an RPC call on all workers.
71
@@ -91,7 +90,7 @@ class ExecutorBase(ABC):
90 """
91 raise NotImplementedError
92
94 - def determine_num_available_blocks(self) -> Tuple[int, int]:
93 + def determine_num_available_blocks(self) -> tuple[int, int]:
94 """Determine the number of available blocks for the GPU KV cache and
95 swappable CPU KV cache.
96
@@ -99,9 +98,10 @@ class ExecutorBase(ABC):
98 ExecutorBase may require modification of the result, e.g. to ensure the
99 selected cache sizes are compatible with all workers.
100
102 - Returns a Tuple[num_gpu_blocks, num_cpu_blocks], where num_gpu_blocks
103 - are blocks that are "active" on the device and can be appended to.
104 - num_cpu_blocks refers to "swapped" blocks in CPU memory and cannot be
101 + Returns a tuple `(num_gpu_blocks, num_cpu_blocks)`, where
102 + `num_gpu_blocks` are blocks that are "active" on the device and can be
103 + appended to.
104 + `num_cpu_blocks` refers to "swapped" blocks in CPU memory and cannot be
105 appended to.
106 """
107 results = self.collective_rpc("determine_num_available_blocks")
@@ -127,16 +127,15 @@ class ExecutorBase(ABC):
127 self.collective_rpc("initialize_cache",
128 args=(num_gpu_blocks, num_cpu_blocks))
129
130 + @deprecated("`llm_engine.model_executor.apply_model` will no longer work "
131 + "in V1 Engine. Please replace with `llm_engine.apply_model` "
132 + "and set `VLLM_ALLOW_INSECURE_SERIALIZATION=1`.")
133 def apply_model(self, func: Callable[[nn.Module], _R]) -> list[_R]:
134 """
135 Run a function directly on the model inside each worker,
136 returning the result for each of them.
137 """
135 -
136 - def rpc_func(worker: WorkerBase) -> _R:
137 - return func(worker.get_model())
138 -
139 - return self.collective_rpc(rpc_func)
138 + return self.collective_rpc("apply_model", args=(func, ))
139
140 @cached_property # Avoid unnecessary RPC calls
141 def supported_tasks(self) -> tuple[SupportedTask, ...]:
@@ -308,8 +307,8 @@ class DistributedExecutorBase(ExecutorBase):
307 def collective_rpc(self,
308 method: Union[str, Callable],
309 timeout: Optional[float] = None,
311 - args: Tuple = (),
312 - kwargs: Optional[Dict] = None) -> List[Any]:
310 + args: tuple = (),
311 + kwargs: Optional[dict[str, Any]] = None) -> list[Any]:
312 return self._run_workers(method, *args, **(kwargs or {}))
313
314 @abstractmethod
vllm/v1/engine/llm_engine.py
+6 -1
@@ -5,6 +5,7 @@ from collections.abc import Mapping
5 from copy import copy
6 from typing import Any, Callable, Optional, Union
7
8 +import torch.nn as nn
9 from typing_extensions import TypeVar
10
11 import vllm.envs as envs
@@ -33,6 +34,7 @@ from vllm.v1.metrics.loggers import (PrometheusStatLogger, StatLoggerBase,
34 StatLoggerFactory)
35 from vllm.v1.metrics.reader import Metric, get_metrics_snapshot
36 from vllm.v1.metrics.stats import IterationStats
37 +from vllm.v1.worker.worker_base import WorkerBase
38
39 logger = init_logger(__name__)
40
@@ -319,12 +321,15 @@ class LLMEngine:
321 return self.engine_core.pin_lora(lora_id)
322
323 def collective_rpc(self,
322 - method: Union[str, Callable[..., _R]],
324 + method: Union[str, Callable[[WorkerBase], _R]],
325 timeout: Optional[float] = None,
326 args: tuple = (),
327 kwargs: Optional[dict[str, Any]] = None) -> list[_R]:
328 return self.engine_core.collective_rpc(method, timeout, args, kwargs)
329
330 + def apply_model(self, func: Callable[[nn.Module], _R]) -> list[_R]:
331 + return self.collective_rpc("apply_model", args=(func, ))
332 +
333 def __del__(self):
334 if dp_group := getattr(self, "dp_group", None):
335 stateless_destroy_torch_distributed_process_group(dp_group)
vllm/worker/worker_base.py
+8 -1
@@ -5,7 +5,8 @@ import dataclasses
5 import os
6 import time
7 from abc import abstractmethod
8 -from typing import Any, Dict, List, Optional, Set, Tuple, Type, Union
8 +from typing import (Any, Callable, Dict, List, Optional, Set, Tuple, Type,
9 + TypeVar, Union)
10
11 import cloudpickle
12 import torch
@@ -28,6 +29,8 @@ from vllm.worker.model_runner_base import (BroadcastableModelInput,
29
30 logger = init_logger(__name__)
31
32 +_R = TypeVar("_R")
33 +
34
35 @warn_for_unimplemented_methods
36 class WorkerBase:
@@ -70,6 +73,10 @@ class WorkerBase:
73 def get_model(self) -> nn.Module:
74 raise NotImplementedError
75
76 + def apply_model(self, fn: Callable[[nn.Module], _R]) -> _R:
77 + """Apply a function on the model inside this worker."""
78 + return fn(self.get_model())
79 +
80 def load_model(self) -> None:
81 """Load model onto target device."""
82 raise NotImplementedError