From a11739b840f1e471d9b4fc5ab03abfa468179ecd Mon Sep 17 00:00:00 2001 From: theissenhelen Date: Thu, 3 Oct 2024 09:29:44 +0000 Subject: [PATCH 1/4] feat: make flash-attn configurabele --- src/anemoi/training/config/model/transformer.yaml | 1 + 1 file changed, 1 insertion(+) diff --git a/src/anemoi/training/config/model/transformer.yaml b/src/anemoi/training/config/model/transformer.yaml index b26c9ecc..6654c596 100644 --- a/src/anemoi/training/config/model/transformer.yaml +++ b/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 + use_flash_attention: True # Transformer encoder: _target_: anemoi.models.layers.mapper.GraphTransformerForwardMapper From a2ed7512e22d8f426152e9c26560b9fefdd8e137 Mon Sep 17 00:00:00 2001 From: theissenhelen Date: Thu, 3 Oct 2024 11:28:32 +0000 Subject: [PATCH 2/4] docs: update changelog --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 80a4bcca..cdf03138 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -137,6 +137,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 From dbd6a8e66b4e09d4217f1b28fc749d855c38735a Mon Sep 17 00:00:00 2001 From: theissenhelen Date: Wed, 27 Nov 2024 12:34:15 +0000 Subject: [PATCH 3/4] fix: adjust config to models implementation --- src/anemoi/training/config/model/transformer.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/anemoi/training/config/model/transformer.yaml b/src/anemoi/training/config/model/transformer.yaml index 6654c596..c7ad0e71 100644 --- a/src/anemoi/training/config/model/transformer.yaml +++ b/src/anemoi/training/config/model/transformer.yaml @@ -14,7 +14,7 @@ processor: num_heads: 16 # GraphTransformer or Transformer only window_size: 512 dropout_p: 0.0 # GraphTransformer - use_flash_attention: True # Transformer + attention_implementation: flash attention # Transformer encoder: _target_: anemoi.models.layers.mapper.GraphTransformerForwardMapper From 6064af381e9fa2ba199f59b94f377f09a3ff37b5 Mon Sep 17 00:00:00 2001 From: theissenhelen Date: Thu, 19 Dec 2024 13:51:01 +0000 Subject: [PATCH 4/4] fix: typo --- src/anemoi/training/config/model/transformer.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/anemoi/training/config/model/transformer.yaml b/src/anemoi/training/config/model/transformer.yaml index 69f01d9a..ef56d36e 100644 --- a/src/anemoi/training/config/model/transformer.yaml +++ b/src/anemoi/training/config/model/transformer.yaml @@ -14,7 +14,7 @@ processor: num_heads: 16 # GraphTransformer or Transformer only window_size: 512 dropout_p: 0.0 # GraphTransformer - attention_implementation: flash attention # Transformer + attention_implementation: flash_attention # Transformer encoder: _target_: anemoi.models.layers.mapper.GraphTransformerForwardMapper