feat: make alibi_slope cinfigurable in block, chunk processor

ecmwf · Oct 1, 2024 · 91533c6 · 91533c6
1 parent bbd89dc
commit 91533c6
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 4 deletions.
diff --git a/src/anemoi/models/layers/block.py b/src/anemoi/models/layers/block.py
@@ -64,7 +64,8 @@ def __init__(
         window_size: int,
         dropout_p: float = 0.0,
         use_flash_attention: bool = False,
-        softcap: float = 0.0,
+        softcap: float | None = 0.0,
+        alibi_slopes: Tensor | None = None,
     ):
         super().__init__()
 
@@ -85,6 +86,7 @@ def __init__(
             dropout_p=dropout_p,
             use_flash_attention=use_flash_attention,
             softcap=softcap,
+            alibi_slopes=alibi_slopes,
         )
 
         self.mlp = nn.Sequential(

diff --git a/src/anemoi/models/layers/chunk.py b/src/anemoi/models/layers/chunk.py
@@ -75,7 +75,8 @@ def __init__(
         activation: str = "GELU",
         dropout_p: float = 0.0,
         use_flash_attention: bool = False,
-        softcap: float = 0.0,
+        softcap: float | None = 0.0,
+        alibi_slopes: Tensor | None = None,
     ) -> None:
         """Initialize TransformerProcessor.
 
@@ -104,8 +105,9 @@ def __init__(
             activation=activation,
             window_size=window_size,
             dropout_p=dropout_p,
-            softcap=softcap,
             use_flash_attention=use_flash_attention,
+            softcap=softcap,
+            alibi_slopes=alibi_slopes,
         )
 
     def forward(

diff --git a/src/anemoi/models/layers/processor.py b/src/anemoi/models/layers/processor.py
@@ -97,7 +97,8 @@ def __init__(
         mlp_hidden_ratio: int = 4,
         dropout_p: float = 0.1,
         use_flash_attention: bool = False,
-        softcap: float = 0.0,
+        softcap: float | None = 0.0,
+        alibi_slopes: Tensor | None = None,
         **kwargs,
     ) -> None:
         """Initialize TransformerProcessor.
@@ -141,6 +142,7 @@ def __init__(
             dropout_p=dropout_p,
             use_flash_attention=use_flash_attention,
             softcap=softcap,
+            alibi_slopes=alibi_slopes,
         )
 
         self.offload_layers(cpu_offload)