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 15, 17:15 UTC
[Bugfix] [torch.compile] Add Dynamo metrics context during compilation (#15639)
Signed-off-by: luka <luka@neuralmagic.com>
Luka Govedič committed
Mar 28, 2025 at 16:01 UTC
04437e313dbbf4427734a3ed2d1d650efc57ef66
2 files changed
+90
-29
tests/compile/test_full_graph.py
+54
-27
@@ -2,21 +2,20 @@
2
3
from __future__ import annotations
4
5
-from typing import Any
5
+from typing import Any, Union
6
7
import pytest
8
import torch
9
10
from tests.quantization.utils import is_quant_method_supported
11
from vllm import LLM, SamplingParams
12
-from vllm.config import CompilationLevel
12
+from vllm.config import CompilationConfig, CompilationLevel
13
from vllm.platforms import current_platform
14
15
from ..utils import create_new_process_for_each_test
16
17
18
-@pytest.fixture(params=None, name="model_info")
19
-def models_list_fixture(request):
18
+def models_list(all: bool):
19
TEST_MODELS: list[tuple[str, dict[str, Any]]] = [
20
("facebook/opt-125m", {}),
21
("nm-testing/tinyllama-oneshot-w8w8-test-static-shape-change", {
@@ -33,6 +32,9 @@ def models_list_fixture(request):
32
("meta-llama/Llama-3.2-1B-Instruct", {}),
33
]
34
35
+ if not all:
36
+ return TEST_MODELS
37
+
38
if is_quant_method_supported("aqlm"):
39
TEST_MODELS.append(("ISTA-DASLab/Llama-2-7b-AQLM-2Bit-1x16-hf", {
40
"quantization": "aqlm"
@@ -77,7 +79,7 @@ def models_list_fixture(request):
79
"optimization_level",
80
[CompilationLevel.DYNAMO_ONCE, CompilationLevel.PIECEWISE],
81
)
80
-@pytest.mark.parametrize("model_info", "", indirect=True)
82
+@pytest.mark.parametrize("model_info", models_list(all=True))
83
@create_new_process_for_each_test()
84
def test_full_graph(
85
monkeypatch: pytest.MonkeyPatch,
@@ -91,25 +93,50 @@ def test_full_graph(
93
m.setenv("VLLM_TEST_DYNAMO_FULLGRAPH_CAPTURE", "1")
94
print(f"MODEL={model}")
95
94
- prompts = [
95
- "Hello, my name is",
96
- "The president of the United States is",
97
- "The capital of France is",
98
- "The future of AI is",
99
- ]
100
- sampling_params = SamplingParams(temperature=0)
101
- llm = LLM(
102
- model=model,
103
- enforce_eager=True,
104
- tensor_parallel_size=1,
105
- disable_custom_all_reduce=True,
106
- compilation_config=optimization_level,
107
- **model_kwargs,
108
- )
109
- outputs = llm.generate(prompts, sampling_params)
110
-
111
- # Print the outputs.
112
- for output in outputs:
113
- prompt = output.prompt
114
- generated_text = output.outputs[0].text
115
- print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
96
+ run_model(optimization_level, model, model_kwargs)
97
+
98
+
99
+# TODO(luka) add other supported compilation config scenarios here
100
+@pytest.mark.parametrize(
101
+ "compilation_config",
102
+ # additional compile sizes
103
+ [
104
+ CompilationConfig(level=CompilationLevel.PIECEWISE,
105
+ compile_sizes=[1, 2])
106
+ ])
107
+# only test some of the models
108
+@pytest.mark.parametrize("model_info", models_list(all=False))
109
+@create_new_process_for_each_test()
110
+def test_custom_compile_config(
111
+ model_info: tuple[str, dict[str, Any]],
112
+ compilation_config: CompilationConfig,
113
+):
114
+ model, model_kwargs = model_info
115
+ print(f"MODEL={model}")
116
+ run_model(compilation_config, model, model_kwargs)
117
+
118
+
119
+def run_model(compile_config: Union[int, CompilationConfig], model: str,
120
+ model_kwargs: dict[str, Any]):
121
+ prompts = [
122
+ "Hello, my name is",
123
+ "The president of the United States is",
124
+ "The capital of France is",
125
+ "The future of AI is",
126
+ ]
127
+ sampling_params = SamplingParams(temperature=0)
128
+ llm = LLM(
129
+ model=model,
130
+ enforce_eager=True,
131
+ tensor_parallel_size=1,
132
+ disable_custom_all_reduce=True,
133
+ compilation_config=compile_config,
134
+ **model_kwargs,
135
+ )
136
+ outputs = llm.generate(prompts, sampling_params)
137
+
138
+ # Print the outputs.
139
+ for output in outputs:
140
+ prompt = output.prompt
141
+ generated_text = output.outputs[0].text
142
+ print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
vllm/compilation/compiler_interface.py
+36
-2
@@ -1,6 +1,8 @@
1
# SPDX-License-Identifier: Apache-2.0
2
+import contextlib
3
import copy
4
import hashlib
5
+import importlib.metadata
6
import os
7
from contextlib import ExitStack
8
from typing import Any, Callable, Dict, List, Optional, Tuple
@@ -9,6 +11,7 @@ from unittest.mock import patch
11
import torch
12
import torch._inductor.compile_fx
13
import torch.fx as fx
14
+from packaging.version import Version
15
16
from vllm.config import VllmConfig
17
@@ -285,6 +288,9 @@ class InductorAdaptor(CompilerInterface):
288
"torch._inductor.codecache.FxGraphCache._check_can_cache",
289
_check_can_cache))
290
291
+ # Dynamo metrics context, see method for more details.
292
+ stack.enter_context(self.metrics_context())
293
+
294
compiled_graph = compile_fx(
295
graph,
296
example_inputs,
@@ -309,8 +315,14 @@ class InductorAdaptor(CompilerInterface):
315
hash_str = handle[0]
316
317
from torch._inductor.codecache import FxGraphCache
312
- with patch("torch._inductor.codecache.FxGraphCache._get_shape_env",
313
- lambda *args, **kwargs: AlwaysHitShapeEnv()):
318
+ with ExitStack() as exit_stack:
319
+ exit_stack.enter_context(
320
+ patch("torch._inductor.codecache.FxGraphCache._get_shape_env",
321
+ lambda *args, **kwargs: AlwaysHitShapeEnv()))
322
+
323
+ # Dynamo metrics context, see method for more details.
324
+ exit_stack.enter_context(self.metrics_context())
325
+
326
if torch.__version__.startswith("2.5"):
327
inductor_compiled_graph = FxGraphCache._lookup_graph(
328
hash_str, example_inputs, True, False)
@@ -351,6 +363,28 @@ class InductorAdaptor(CompilerInterface):
363
364
return compiled_graph
365
366
+ def metrics_context(self) -> contextlib.AbstractContextManager:
367
+ """
368
+ This method returns the Dynamo metrics context (if it exists,
369
+ otherwise a null context). It is used by various compile components.
370
+ Present in torch>=2.6, it's used inside FxGraphCache in
371
+ torch==2.6 (but not after). It might also be used in various other
372
+ torch.compile internal functions.
373
+
374
+ Because it is re-entrant, we always set it (even if entering via Dynamo
375
+ and the context was already entered). We might want to revisit if it
376
+ should be set at a different level of compilation.
377
+
378
+ This is likely a bug in PyTorch: public APIs should not rely on
379
+ manually setting up internal contexts. But we also rely on non-public
380
+ APIs which might not provide these guarantees.
381
+ """
382
+ if Version(importlib.metadata.version('torch')) >= Version("2.6"):
383
+ import torch._dynamo.utils
384
+ return torch._dynamo.utils.get_metrics_context()
385
+ else:
386
+ return contextlib.nullcontext()
387
+
388
389
class EagerAdaptor(CompilerInterface):
390
name = "eager"