feat: make flash-attn configurabele

CHANGELOG.md

@@ -152,6 +152,7 @@ Keep it human-readable, your future self will thank you!
 - Remove credential prompt from mlflow login, replace with seed refresh token via web - [#78](https://github.com/ecmwf/anemoi-training/pull/78)
 - Update CODEOWNERS
 - Change how mlflow measures CPU Memory usage - [94](https://github.com/ecmwf/anemoi-training/pull/94)
+- Updated configuration examples to configure flash-attn - [#73] (https://github.com/ecmwf/anemoi-training/pull/73)
 
 ## [0.1.0 - Anemoi training - First release](https://github.com/ecmwf/anemoi-training/releases/tag/0.1.0) - 2024-08-16
 

src/anemoi/training/config/model/transformer.yaml

@@ -14,6 +14,7 @@ processor:
   num_heads: 16 # GraphTransformer or Transformer only
   window_size: 512
   dropout_p: 0.0 # GraphTransformer
+  attention_implementation: flash_attention # Transformer
 
 encoder:
   _target_: anemoi.models.layers.mapper.GraphTransformerForwardMapper