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

Loosen requirement on packaging dependency (#1758)

上级 dac9056e
......@@ -3,7 +3,11 @@ Copyright 2020 The Microsoft DeepSpeed Team
"""
import warnings
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):
......@@ -62,8 +66,14 @@ class SparseAttnBuilder(OpBuilder):
f"please install triton==1.0.0 if you want to use sparse attention")
return False
if pkg_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(
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.
先完成此消息的编辑!
想要评论请 注册