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, 14:45 UTC
Fix the torch version parsing logic (#15857)
Lu Fang committed
Apr 10, 2025 at 07:37 UTC
7678fcd5b6d64084aeabeb17251d388e251ba4c9
4 files changed
+26
-11
vllm/compilation/compiler_interface.py
+2
-3
@@ -2,7 +2,6 @@
2
import contextlib
3
import copy
4
import hashlib
5
-import importlib.metadata
5
import os
6
from contextlib import ExitStack
7
from typing import Any, Callable, Dict, List, Optional, Tuple
@@ -11,9 +10,9 @@ from unittest.mock import patch
10
import torch
11
import torch._inductor.compile_fx
12
import torch.fx as fx
14
-from packaging.version import Version
13
14
from vllm.config import VllmConfig
15
+from vllm.utils import is_torch_equal_or_newer
16
17
18
class CompilerInterface:
@@ -379,7 +378,7 @@ class InductorAdaptor(CompilerInterface):
378
manually setting up internal contexts. But we also rely on non-public
379
APIs which might not provide these guarantees.
380
"""
382
- if Version(importlib.metadata.version('torch')) >= Version("2.6"):
381
+ if is_torch_equal_or_newer("2.6"):
382
import torch._dynamo.utils
383
return torch._dynamo.utils.get_metrics_context()
384
else:
vllm/compilation/inductor_pass.py
+3
-3
@@ -1,17 +1,17 @@
1
# SPDX-License-Identifier: Apache-2.0
2
3
import hashlib
4
-import importlib.metadata
4
import inspect
5
import json
6
import types
7
from typing import Any, Callable, Dict, Optional, Union
8
9
import torch
11
-from packaging.version import Version
10
from torch import fx
11
14
-if Version(importlib.metadata.version('torch')) >= Version("2.6"):
12
+from vllm.utils import is_torch_equal_or_newer
13
+
14
+if is_torch_equal_or_newer("2.6"):
15
from torch._inductor.custom_graph_pass import CustomGraphPass
16
else:
17
# CustomGraphPass is not present in 2.5 or lower, import our version
vllm/config.py
+3
-5
@@ -4,7 +4,6 @@ import ast
4
import copy
5
import enum
6
import hashlib
7
-import importlib.metadata
7
import json
8
import sys
9
import warnings
@@ -18,7 +17,6 @@ from typing import (TYPE_CHECKING, Any, Callable, ClassVar, Final, Literal,
17
Optional, Protocol, Union)
18
19
import torch
21
-from packaging.version import Version
20
from pydantic import BaseModel, Field, PrivateAttr
21
from torch.distributed import ProcessGroup, ReduceOp
22
from transformers import PretrainedConfig
@@ -40,8 +38,8 @@ from vllm.transformers_utils.config import (
38
from vllm.transformers_utils.s3_utils import S3Model
39
from vllm.transformers_utils.utils import is_s3, maybe_model_redirect
40
from vllm.utils import (GiB_bytes, LayerBlockType, cuda_device_count_stateless,
43
- get_cpu_memory, get_open_port, random_uuid,
44
- resolve_obj_by_qualname)
41
+ get_cpu_memory, get_open_port, is_torch_equal_or_newer,
42
+ random_uuid, resolve_obj_by_qualname)
43
44
if TYPE_CHECKING:
45
from ray.util.placement_group import PlacementGroup
@@ -3285,7 +3283,7 @@ class CompilationConfig(BaseModel):
3283
# and it is not yet a priority. RFC here:
3284
# https://github.com/vllm-project/vllm/issues/14703
3285
3288
- if Version(importlib.metadata.version('torch')) >= Version("2.6"):
3286
+ if is_torch_equal_or_newer("2.6"):
3287
KEY = 'enable_auto_functionalized_v2'
3288
if KEY not in self.inductor_compile_config:
3289
self.inductor_compile_config[KEY] = False
vllm/utils.py
+18
@@ -53,6 +53,7 @@ import torch.types
53
import yaml
54
import zmq
55
import zmq.asyncio
56
+from packaging import version
57
from packaging.version import Version
58
from torch.library import Library
59
from typing_extensions import Never, ParamSpec, TypeIs, assert_never
@@ -2580,3 +2581,20 @@ def sha256(input) -> int:
2581
input_bytes = pickle.dumps(input, protocol=pickle.HIGHEST_PROTOCOL)
2582
return int.from_bytes(hashlib.sha256(input_bytes).digest(),
2583
byteorder="big")
2584
+
2585
+
2586
+def is_torch_equal_or_newer(target: str) -> bool:
2587
+ """Check if the installed torch version is >= the target version.
2588
+
2589
+ Args:
2590
+ target: a version string, like "2.6.0".
2591
+
2592
+ Returns:
2593
+ Whether the condition meets.
2594
+ """
2595
+ try:
2596
+ torch_version = version.parse(str(torch.__version__))
2597
+ return torch_version >= version.parse(target)
2598
+ except Exception:
2599
+ # Fallback to PKG-INFO to load the package info, needed by the doc gen.
2600
+ return Version(importlib.metadata.version('torch')) >= Version(target)