add version check

mikecovlee · Jul 18, 2024 · 415c286 · 415c286
1 parent af1457e
commit 415c286
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/mlora/models/modeling_gemma2.py b/mlora/models/modeling_gemma2.py
@@ -26,7 +26,7 @@
 )
 from mlora.models.modeling_gemma import GemmaEmbedding, GemmaRMSNorm
 from mlora.models.modeling_llama import LlamaMLP
-from mlora.utils import copy_parameters
+from mlora.utils import copy_parameters, is_package_available
 
 if is_flash_attn_2_available():
     from flash_attn import flash_attn_func, flash_attn_varlen_func
@@ -36,6 +36,10 @@
         inspect.signature(flash_attn_func).parameters
     )
 
+    assert is_package_available(
+        "flash_attn", "2.6.0"
+    ), "Gemma2 requires flash_attn>=2.6.0"
+
 
 @dataclass
 class Gemma2Config(LLMModelConfig):