diff --git a/tests/compile/fullgraph/test_basic_correctness.py b/tests/compile/fullgraph/test_basic_correctness.py index 859c4844d6..304e6d7cae 100644 --- a/tests/compile/fullgraph/test_basic_correctness.py +++ b/tests/compile/fullgraph/test_basic_correctness.py @@ -56,8 +56,10 @@ class TestSetting: ), ], ) +@pytest.mark.parametrize("use_v2_model_runner", [False, True]) def test_compile_correctness( test_setting: TestSetting, + use_v2_model_runner: bool, ): model = test_setting.model model_args = test_setting.model_args @@ -75,6 +77,10 @@ def test_compile_correctness( f"--attention-backend={attn_backend}", ] + # Pin every compared setting to the same model runner so that only the + # compilation mode differs within a comparison. + runner_env = {"VLLM_USE_V2_MODEL_RUNNER": str(int(use_v2_model_runner))} + all_args: list[list[str]] = [] all_envs: list[dict[str, str] | None] = [] @@ -86,7 +92,7 @@ def test_compile_correctness( CompilationMode.VLLM_COMPILE, ]: all_args.append(final_args + [f"-cc.mode={mode.name}", "-cc.backend=inductor"]) - all_envs.append({}) + all_envs.append(dict(runner_env)) # inductor will change the output, so we only compare if the output # is close, not exactly the same. compare_all_settings( @@ -94,7 +100,6 @@ def test_compile_correctness( all_args, all_envs, method=method if method != "generate" else "generate_close", - force_v1_runner=True, ) all_envs.clear() @@ -108,5 +113,5 @@ def test_compile_correctness( CompilationMode.VLLM_COMPILE, ]: all_args.append(final_args + [f"-cc.mode={mode.name}", "-cc.backend=eager"]) - all_envs.append({}) - compare_all_settings(model, all_args, all_envs, method=method, force_v1_runner=True) + all_envs.append(dict(runner_env)) + compare_all_settings(model, all_args, all_envs, method=method) diff --git a/tests/compile/test_config.py b/tests/compile/test_config.py index 74ed0d34db..24cb816b5d 100644 --- a/tests/compile/test_config.py +++ b/tests/compile/test_config.py @@ -215,7 +215,9 @@ def test_use_cudagraphs( # forked needed to workaround https://github.com/vllm-project/vllm/issues/21073 @pytest.mark.forked -def test_stock_torch_compile(vllm_runner, monkeypatch): +@pytest.mark.parametrize("use_v2_model_runner", [False, True]) +def test_stock_torch_compile(vllm_runner, monkeypatch, use_v2_model_runner): + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", str(int(use_v2_model_runner))) # Disable multiprocessing so that the counter is in the same process monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") @@ -226,9 +228,10 @@ def test_stock_torch_compile(vllm_runner, monkeypatch): "facebook/opt-125m", compilation_config={"mode": CompilationMode.STOCK_TORCH_COMPILE}, gpu_memory_utilization=0.4, - ) as _, + ) as runner, ): - pass + outputs = runner.generate_greedy(["Hello, my name is"], max_tokens=5) + assert outputs[0][0] # forked needed to workaround https://github.com/vllm-project/vllm/issues/21073 diff --git a/tests/compile/test_wrapper.py b/tests/compile/test_wrapper.py index 5e0755ff71..45dc148cda 100644 --- a/tests/compile/test_wrapper.py +++ b/tests/compile/test_wrapper.py @@ -7,7 +7,11 @@ import os import pytest import torch -from vllm.compilation.wrapper import TorchCompileWithNoGuardsWrapper +from vllm.compilation.counter import compilation_counter +from vllm.compilation.wrapper import ( + TorchCompileWithNoGuardsWrapper, + compile_model_with_stock_torch, +) from vllm.config import ( CompilationConfig, CompilationMode, @@ -117,6 +121,44 @@ def test_torch_compile_wrapper(use_bytecode_hook, monkeypatch): raise AssertionError("expected an exception to be raised") +def test_compile_model_with_stock_torch(monkeypatch): + """The whole model is compiled in place with the configured backend.""" + graphs: list[torch.fx.GraphModule] = [] + + def recording_backend(gm: torch.fx.GraphModule, example_inputs): + graphs.append(gm) + return gm.forward + + monkeypatch.setattr( + CompilationConfig, + "init_backend", + lambda self, vllm_config, *args, **kwargs: recording_backend, + ) + + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(4, 2) + + def forward(self, x: torch.Tensor): + return self.linear(x).relu() + + vllm_config = VllmConfig() + vllm_config.compilation_config = CompilationConfig( + mode=CompilationMode.STOCK_TORCH_COMPILE + ) + model = Model() + x = torch.randn(3, 4) + expected = model(x) + + torch._dynamo.reset() + with compilation_counter.expect(stock_torch_compile_count=1): + compile_model_with_stock_torch(model, vllm_config) + assert not graphs # compilation is deferred to the first call + torch.testing.assert_close(model(x), expected) + assert len(graphs) == 1 + + if __name__ == "__main__": # Run with both parameter values diff --git a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py index 5a32aa63d8..7c759894fe 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py +++ b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py @@ -7,6 +7,7 @@ from typing import Any import torch +from vllm.config import CompilationConfig, CompilationMode from vllm.model_executor.warmup.jit_warmup import JitWarmupRegistry from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT from vllm.v1.worker.gpu import eplb_utils as eplb @@ -68,7 +69,9 @@ def _make_runner(**overrides: Any) -> Any: enable_batch_sharded_sampling=False, eplb_config=SimpleNamespace(log_balancedness=True), ) + runner.compilation_config = CompilationConfig(mode=CompilationMode.NONE) runner.vllm_config = SimpleNamespace( + compilation_config=runner.compilation_config, load_config=runner.load_config, model_config=runner.model_config, ) diff --git a/vllm/compilation/wrapper.py b/vllm/compilation/wrapper.py index 1df30a70a8..78b34f8959 100644 --- a/vllm/compilation/wrapper.py +++ b/vllm/compilation/wrapper.py @@ -12,7 +12,13 @@ from typing import Any, ParamSpec, TypeVar import torch import vllm.envs as envs -from vllm.config import CompilationMode, CUDAGraphMode, get_current_vllm_config +from vllm.compilation.counter import compilation_counter +from vllm.config import ( + CompilationMode, + CUDAGraphMode, + VllmConfig, + get_current_vllm_config, +) from vllm.config.compilation import DynamicShapesType from vllm.logger import init_logger from vllm.utils.nvtx_pytorch_hooks import layerwise_nvtx_marker_context @@ -44,6 +50,17 @@ def _compilation_context() -> Generator[None, None, None]: torch._dynamo.config.accumulated_cache_size_limit = original_accumulated_cache +def compile_model_with_stock_torch( + model: torch.nn.Module, vllm_config: VllmConfig +) -> None: + from vllm.env_override import _apply_constrain_to_fx_strides_patch + + _apply_constrain_to_fx_strides_patch() + backend = vllm_config.compilation_config.init_backend(vllm_config) + compilation_counter.stock_torch_compile_count += 1 + model.compile(fullgraph=True, backend=backend) + + class TorchCompileWithNoGuardsWrapper: """A wrapper class for torch.compile, it ensures that all guards are dropped when CompilationMode is not CompilationMode.STOCK_TORCH_COMPILE. diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 1b67d0e631..8022cd0b88 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -2041,6 +2041,7 @@ class VllmConfig: self._validate_mm_processor_device() if self.use_v2_model_runner: + self._disable_cudagraphs_for_v2_stock_torch_compile() self._validate_v2_model_runner() else: self._validate_v1_model_runner() @@ -3045,9 +3046,6 @@ class VllmConfig: unsupported: list[str] = [] speculative_config = self.speculative_config - if self.compilation_config.mode == CompilationMode.STOCK_TORCH_COMPILE: - unsupported.append("stock torch.compile") - if ( self.compilation_config.pass_config.enable_sp and self.parallel_config.tensor_parallel_size > 1 @@ -3301,6 +3299,28 @@ class VllmConfig: "Enable attribution or use --enforce-eager." ) + def _disable_cudagraphs_for_v2_stock_torch_compile(self) -> None: + """Run stock torch.compile without CUDA graphs in Model Runner V2. + + V1 never wraps a stock-compiled model in CUDAGraphWrapper, so it runs + without CUDA graphs. V2's CUDA graph manager would otherwise capture + FULL graphs around the stock-compiled model, so disable them to match. + """ + compilation_config = self.compilation_config + if ( + compilation_config.mode != CompilationMode.STOCK_TORCH_COMPILE + or compilation_config.cudagraph_mode == CUDAGraphMode.NONE + ): + return + logger.info_once( + "CUDA graphs are not supported with stock torch.compile in Model " + "Runner V2. Overriding cudagraph_mode %s to NONE.", + compilation_config.cudagraph_mode.name, + ) + compilation_config.cudagraph_mode = CUDAGraphMode.NONE + compilation_config.max_cudagraph_capture_size = 0 + compilation_config.cudagraph_capture_sizes = [] + def _validate_v2_model_runner(self) -> None: """Check for features not yet supported by the V2 model runner.""" if not HAS_TRITON: diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 50eff377be..800ae825d4 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -31,8 +31,9 @@ import torch.nn as nn import vllm.envs as envs from vllm.compilation.counter import compilation_counter from vllm.compilation.cuda_graph import CUDAGraphStat +from vllm.compilation.wrapper import compile_model_with_stock_torch from vllm.config import VllmConfig -from vllm.config.compilation import CUDAGraphMode +from vllm.config.compilation import CompilationMode, CUDAGraphMode from vllm.distributed.aux_output_connector.worker import ( AuxOutputWorkerConnector, get_aux_output_connector, @@ -522,6 +523,9 @@ class GPUModelRunner(LoRAModelRunnerMixin): get_offloader().post_init() + if self.compilation_config.mode == CompilationMode.STOCK_TORCH_COMPILE: + compile_model_with_stock_torch(self.model, self.vllm_config) + @property def dp_size(self) -> int: # Not cached: elastic EP rewrites parallel_config in place on reconfigure.