Skip to content

Commit

Permalink
Merge branch 'cherry-pick-fdee903f' into 'core_r0.7.0'
Browse files Browse the repository at this point in the history
Merge branch 'gg_guard' into 'main'

See merge request ADLR/megatron-lm!1516
  • Loading branch information
jaredcasper committed May 30, 2024
2 parents e2627c7 + e24256a commit a967adf
Showing 1 changed file with 11 additions and 1 deletion.
12 changes: 11 additions & 1 deletion megatron/core/transformer/moe/grouped_gemm_util.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,9 @@
# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.

from importlib.metadata import version

from pkg_resources import packaging

try:
import grouped_gemm
except ImportError:
Expand All @@ -13,7 +17,13 @@ def grouped_gemm_is_available():
def assert_grouped_gemm_is_available():
assert grouped_gemm_is_available(), (
"Grouped GEMM is not available. Please run "
"`pip install git+https://github.com/fanshiqing/[email protected]`."
"`pip install git+https://github.com/fanshiqing/[email protected]`."
)

_gg_version = packaging.version.Version(version("grouped_gemm"))
assert _gg_version >= packaging.version.Version("1.1.2"), (
"Grouped GEMM should be v1.1.2 or newer. Please run "
"`pip install git+https://github.com/fanshiqing/[email protected]`."
)


Expand Down

0 comments on commit a967adf

Please sign in to comment.