[MRV2] Support stock torch.compile mode (#59079)

Signed-off-by: cheesecake <farzanaman99@gmail.com>
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>

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)

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

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
 

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,
     )