未验证 提交 dbe8ee16 编写于 作者: J Jeff Rasley 提交者: GitHub

Loosen requirement on packaging dependency (#1758)

上级 dac9056e
...@@ -3,7 +3,11 @@ Copyright 2020 The Microsoft DeepSpeed Team ...@@ -3,7 +3,11 @@ Copyright 2020 The Microsoft DeepSpeed Team
""" """
import warnings import warnings
from .builder import OpBuilder from .builder import OpBuilder
from packaging import version as pkg_version
try:
from packaging import version as pkg_version
except ImportError:
pkg_version = None
class SparseAttnBuilder(OpBuilder): class SparseAttnBuilder(OpBuilder):
...@@ -62,8 +66,14 @@ class SparseAttnBuilder(OpBuilder): ...@@ -62,8 +66,14 @@ class SparseAttnBuilder(OpBuilder):
f"please install triton==1.0.0 if you want to use sparse attention") f"please install triton==1.0.0 if you want to use sparse attention")
return False return False
if pkg_version:
installed_triton = pkg_version.parse(triton.__version__) installed_triton = pkg_version.parse(triton.__version__)
if installed_triton != pkg_version.parse("1.0.0"): triton_mismatch = installed_triton != pkg_version.parse("1.0.0")
else:
installed_triton = triton.__version__
triton_mismatch = installed_triton != "1.0.0"
if triton_mismatch:
self.warning( self.warning(
f"using untested triton version ({installed_triton}), only 1.0.0 is known to be compatible" f"using untested triton version ({installed_triton}), only 1.0.0 is known to be compatible"
) )
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册