Merge branch 'fix_12' into 'core_r0.12.0'

Adding skipif to legacy tests

See merge request ADLR/megatron-lm!3107

Author: ko3n1g

megatron/core/utils.py

@@ -47,6 +47,7 @@
     # This is a WAR for building docs, where torch is not actually imported
     _torch_version = PkgVersion("0.0.0")
 _te_version = None
+_fa_version = None
 
 
 class ExperimentalNotEnabledError(Exception):
@@ -279,6 +280,30 @@ def is_torch_min_version(version, check_equality=True):
     return get_torch_version() > PkgVersion(version)
 
 
+def get_fa_version():
+    """Get Flash attention version from __version__; if not available use pip's. Use caching."""
+
+    def get_fa_version_str():
+        import flash_attn as fa
+
+        if hasattr(fa, '__version__'):
+            return str(fa.__version__)
+        else:
+            return version("flash-attn")
+
+    global _fa_version
+    if _fa_version is None:
+        _fa_version = PkgVersion(get_fa_version_str())
+    return _fa_version
+
+
+def is_fa_min_version(version, check_equality=True):
+    """Check if minimum version of `flash-attn` is installed."""
+    if check_equality:
+        return get_fa_version() >= PkgVersion(version)
+    return get_fa_version() > PkgVersion(version)
+
+
 def ensure_divisibility(numerator, denominator):
     """Ensure that numerator is divisible by the denominator."""
     assert numerator % denominator == 0, "{} is not divisible by {}".format(numerator, denominator)

tests/unit_tests/inference/engines/test_dynamic_engine.py

@@ -5,6 +5,7 @@
 from dataclasses import dataclass
 from typing import List, Optional
 
+import pytest
 import torch
 from tqdm import tqdm
 
@@ -30,6 +31,7 @@
 from megatron.core.models.gpt.gpt_model import GPTModel
 from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
 from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.utils import is_fa_min_version
 from tests.unit_tests.test_utilities import Utils
 
 DynamicInferenceContext.ROUNDER = 4  # decreased from 64 for unit tests.
@@ -310,6 +312,9 @@ def setup_method(self, method):
     def teardown_method(self, method):
         Utils.destroy_model_parallel()
 
+    @pytest.mark.skipif(
+        not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching"
+    )
     def test_simple(self) -> None:
         """Simple test that runs without errors, and validates output."""
 
@@ -336,6 +341,9 @@ def test_simple(self) -> None:
         for request, expected_output in zip(env.requests, expected_outputs):
             assert request.output == expected_output
 
+    @pytest.mark.skipif(
+        not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching"
+    )
     def test_overflow_factor(self) -> None:
         """Test overflow factor arg."""
 
@@ -350,6 +358,9 @@ def test_overflow_factor(self) -> None:
         assert env.engine.context.max_requests == 1120
         assert env.engine.context.max_tokens == 1120
 
+    @pytest.mark.skipif(
+        not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching"
+    )
     def test_request_overflow(self) -> None:
         """Test request overflow."""
         try:
@@ -358,6 +369,9 @@ def test_request_overflow(self) -> None:
             return
         raise Exception("failed.")
 
+    @pytest.mark.skipif(
+        not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching"
+    )
     def test_token_overflow(self) -> None:
         """Test token overflow."""
         try:
@@ -366,6 +380,9 @@ def test_token_overflow(self) -> None:
             return
         raise Exception("failed.")
 
+    @pytest.mark.skipif(
+        not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching"
+    )
     def test_chunk_overflow(self) -> None:
         """Test chunk overflow."""
         env = self._build_test_env(TestConfig())
@@ -378,10 +395,16 @@ def test_chunk_overflow(self) -> None:
             return
         raise Exception("failed.")
 
+    @pytest.mark.skipif(
+        not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching"
+    )
     def test_multi_add(self) -> None:
         """Test adding multiple requests simultaneously."""
         self._run_test(num_gap_steps=0)
 
+    @pytest.mark.skipif(
+        not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching"
+    )
     def test_fixed_output_lengths(self) -> None:
         """Test generating a fixed number of output tokens."""
         self._run_test(use_fixed_output_lengths=True)