test: adjust tests

ecmwf · Nov 27, 2024 · f1be563 · f1be563
1 parent a703688
commit f1be563
Show file tree

Hide file tree

Showing 5 changed files with 46 additions and 9 deletions.
diff --git a/src/anemoi/models/layers/processor.py b/src/anemoi/models/layers/processor.py
@@ -97,7 +97,7 @@ def __init__(
         num_heads: int = 16,
         mlp_hidden_ratio: int = 4,
         dropout_p: float = 0.1,
-        attention_implementation: str = "Flex Attention",
+        attention_implementation: str = "flex attention",
         softcap: float = 0.0,
         use_alibi_slopes: bool = None,
         **kwargs,

diff --git a/tests/layers/block/test_block_transformer.py b/tests/layers/block/test_block_transformer.py
@@ -39,7 +39,14 @@ class TestTransformerProcessorBlock:
     def test_init(self, factor_attention_heads, hidden_dim, num_heads, activation, window_size, dropout_p, softcap):
         num_channels = num_heads * factor_attention_heads
         block = TransformerProcessorBlock(
-            num_channels, hidden_dim, num_heads, activation, window_size, dropout_p=dropout_p, softcap=softcap
+            num_channels,
+            hidden_dim,
+            num_heads,
+            activation,
+            window_size,
+            dropout_p=dropout_p,
+            attention_implementation="scaled dot product attention",
+            softcap=softcap,
         )
         assert isinstance(block, TransformerProcessorBlock)
 
@@ -74,7 +81,14 @@ def test_forward_output(
     ):
         num_channels = num_heads * factor_attention_heads
         block = TransformerProcessorBlock(
-            num_channels, hidden_dim, num_heads, activation, window_size, dropout_p=dropout_p, softcap=softcap
+            num_channels,
+            hidden_dim,
+            num_heads,
+            activation,
+            window_size,
+            dropout_p=dropout_p,
+            attention_implementation="scaled dot product attention",
+            softcap=softcap,
         )
 
         x = torch.randn((batch_size, num_channels))  # .to(torch.float16, non_blocking=True)

diff --git a/tests/layers/chunk/test_chunk_transformer.py b/tests/layers/chunk/test_chunk_transformer.py
@@ -24,6 +24,7 @@ def init(self):
         activation: str = "GELU"
         window_size: int = 13
         dropout_p: float = 0.1
+        attention_implementation = "scaled dot product attention"
 
         # num_heads must be evenly divisible by num_channels for MHSA
         return (
@@ -34,6 +35,7 @@ def init(self):
             activation,
             window_size,
             dropout_p,
+            attention_implementation,
         )
 
     @pytest.fixture
@@ -46,6 +48,7 @@ def processor_chunk(self, init):
             activation,
             window_size,
             dropout_p,
+            attention_implementation,
         ) = init
         return TransformerProcessorChunk(
             num_channels=num_channels,
@@ -55,6 +58,7 @@ def processor_chunk(self, init):
             activation=activation,
             window_size=window_size,
             dropout_p=dropout_p,
+            attention_implementation=attention_implementation,
         )
 
     def test_all_blocks(self, processor_chunk):

diff --git a/tests/layers/processor/test_transformer_processor.py b/tests/layers/processor/test_transformer_processor.py
@@ -26,6 +26,7 @@ def transformer_processor_init():
     mlp_hidden_ratio = 4
     dropout_p = 0.1
     softcap = 0.5
+    attention_implementation = "scaled dot product attention"
     return (
         num_layers,
         window_size,
@@ -37,6 +38,7 @@ def transformer_processor_init():
         mlp_hidden_ratio,
         dropout_p,
         softcap,
+        attention_implementation,
     )
 
 
@@ -53,6 +55,7 @@ def transformer_processor(transformer_processor_init):
         mlp_hidden_ratio,
         dropout_p,
         softcap,
+        attention_implementation,
     ) = transformer_processor_init
     return TransformerProcessor(
         num_layers=num_layers,
@@ -64,6 +67,7 @@ def transformer_processor(transformer_processor_init):
         num_heads=num_heads,
         mlp_hidden_ratio=mlp_hidden_ratio,
         dropout_p=dropout_p,
+        attention_implementation=attention_implementation,
         softcap=softcap,
     )
 
@@ -79,6 +83,7 @@ def test_transformer_processor_init(transformer_processor, transformer_processor
         _num_heads,
         _mlp_hidden_ratio,
         _dropout_p,
+        _attention_implementation,
         _softcap,
     ) = transformer_processor_init
     assert isinstance(transformer_processor, TransformerProcessor)
@@ -98,6 +103,7 @@ def test_transformer_processor_forward(transformer_processor, transformer_proces
         _num_heads,
         _mlp_hidden_ratio,
         _dropout_p,
+        _attention_implementation,
         _softcap,
     ) = transformer_processor_init
     gridsize = 100

diff --git a/tests/layers/test_attention.py b/tests/layers/test_attention.py
@@ -23,12 +23,15 @@
     embed_dim_multiplier=st.integers(min_value=1, max_value=10),
     dropout_p=st.floats(min_value=0.0, max_value=1.0),
     softcap=st.floats(min_value=0.0, max_value=1.0),
+    attention_implementation=st.sampled_from(["scaled dot product attention", "flex attention"]),
 )
-def test_multi_head_self_attention_init(num_heads, embed_dim_multiplier, dropout_p, softcap):
+def test_multi_head_self_attention_init(num_heads, embed_dim_multiplier, dropout_p, softcap, attention_implementation):
     embed_dim = (
         num_heads * embed_dim_multiplier
     )  # TODO: Make assert in MHSA to check if embed_dim is divisible by num_heads
-    mhsa = MultiHeadSelfAttention(num_heads, embed_dim, dropout_p=dropout_p, use_flash_attention=False, softcap=softcap)
+    mhsa = MultiHeadSelfAttention(
+        num_heads, embed_dim, dropout_p=dropout_p, attention_implementation=attention_implementation, softcap=softcap
+    )
 
     assert isinstance(mhsa, nn.Module)
     assert mhsa.num_heads == num_heads
@@ -44,11 +47,16 @@ def test_multi_head_self_attention_init(num_heads, embed_dim_multiplier, dropout
     embed_dim_multiplier=st.integers(min_value=1, max_value=10),
     dropout_p=st.floats(min_value=0.0, max_value=1.0),
     softcap=st.floats(min_value=0.0, max_value=1.0),
+    attention_implementation=st.sampled_from(["scaled dot product attention"]),
 )
 @settings(deadline=None)
-def test_multi_head_self_attention_forward(batch_size, num_heads, embed_dim_multiplier, dropout_p, softcap):
+def test_multi_head_self_attention_forward(
+    batch_size, num_heads, embed_dim_multiplier, dropout_p, softcap, attention_implementation
+):
     embed_dim = num_heads * embed_dim_multiplier
-    mhsa = MultiHeadSelfAttention(num_heads, embed_dim, dropout_p=dropout_p, use_flash_attention=False, softcap=softcap)
+    mhsa = MultiHeadSelfAttention(
+        num_heads, embed_dim, dropout_p=dropout_p, attention_implementation=attention_implementation, softcap=softcap
+    )
 
     x = torch.randn(batch_size * 2, embed_dim)
     shapes = [list(x.shape)]
@@ -64,10 +72,15 @@ def test_multi_head_self_attention_forward(batch_size, num_heads, embed_dim_mult
     embed_dim_multiplier=st.integers(min_value=1, max_value=10),
     dropout_p=st.floats(min_value=0.0, max_value=1.0),
     softcap=st.floats(min_value=0.0, max_value=1.0),
+    attention_implementation=st.sampled_from(["scaled dot product attention"]),
 )
-def test_multi_head_self_attention_backward(batch_size, num_heads, embed_dim_multiplier, dropout_p, softcap):
+def test_multi_head_self_attention_backward(
+    batch_size, num_heads, embed_dim_multiplier, dropout_p, softcap, attention_implementation
+):
     embed_dim = num_heads * embed_dim_multiplier
-    mhsa = MultiHeadSelfAttention(num_heads, embed_dim, dropout_p=dropout_p, use_flash_attention=False, softcap=softcap)
+    mhsa = MultiHeadSelfAttention(
+        num_heads, embed_dim, dropout_p=dropout_p, attention_implementation=attention_implementation, softcap=softcap
+    )
 
     x = torch.randn(batch_size * 2, embed_dim, requires_grad=True)
     shapes = [list(x.shape)]