Skip to content

Commit e24256a

Browse files
committed
Merge branch 'gg_guard' into 'main'
groupedGEMM correct library version in import guard See merge request ADLR/megatron-lm!1513 (cherry picked from commit fdee903) 54202ec9 groupedGEMM correct library version in import guard 8d4d1d6b add package version check
1 parent e2627c7 commit e24256a

File tree

1 file changed

+11
-1
lines changed

1 file changed

+11
-1
lines changed

megatron/core/transformer/moe/grouped_gemm_util.py

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,9 @@
11
# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
22

3+
from importlib.metadata import version
4+
5+
from pkg_resources import packaging
6+
37
try:
48
import grouped_gemm
59
except ImportError:
@@ -13,7 +17,13 @@ def grouped_gemm_is_available():
1317
def assert_grouped_gemm_is_available():
1418
assert grouped_gemm_is_available(), (
1519
"Grouped GEMM is not available. Please run "
16-
"`pip install git+https://github.com/fanshiqing/[email protected]`."
20+
"`pip install git+https://github.com/fanshiqing/[email protected]`."
21+
)
22+
23+
_gg_version = packaging.version.Version(version("grouped_gemm"))
24+
assert _gg_version >= packaging.version.Version("1.1.2"), (
25+
"Grouped GEMM should be v1.1.2 or newer. Please run "
26+
"`pip install git+https://github.com/fanshiqing/[email protected]`."
1727
)
1828

1929

0 commit comments

Comments
 (0)