..

mosaicml · Nov 18, 2023 · 3d1d022 · 3d1d022
1 parent 42a541d
commit 3d1d022
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/tests/test_model.py b/tests/test_model.py
@@ -565,6 +565,12 @@ def test_sequence_id_based_masking(attention_impl: str, device: str,
         pytest.skip(
             f'dail implementation of rope requires gpu and flash attention 2.')
 
+    if attention_impl == 'flash' and (
+            not is_flash_v2_installed(v2_version='v2.1.2')):
+        pytest.skip(
+            'Using sequence id with flash attention requires flash attention v2.1.2 or higher.'
+        )
+
     composer_device = get_device(device)
 
     hf_config = MPTConfig(