Feature/44 make flash attention configurable #47

theissenhelen · 2024-09-20T14:00:00Z

Current setup:

If flash-attn is available in the environment, the MultiHeadSelfAttention module automatically imports the corresponding attention function. In inference however we do not have that information.

Now:

user specifies whether flash-attn should be used in the model config.
adds configurable parameters (soft cap, aLiBi) for flash attention
for aLiB:i adds a function to compute the slopes according to the number of attention heads

src/anemoi/models/layers/attention.py

* fix: change pre-cmmit autoupdate schedule to monthly * fix: change the merge strategy for Changelog to Union * fix: add .envrc to .gitignore * ci: ignore pre-commit-config and readthedocs for changelog updates * ci: fix to correct hpc workflow call * fix: update precommit config * chore: update pre-commits * feat: add codeowners file * chore: update dependencies * ci: add hpc-config * docs: changelog * fix: respond to review comments --------- Co-authored-by: Jesper Dramsch <[email protected]>

* feat: add configurability to dropout in MultiHeadSelfAttention Co-authored-by: Rilwan (Akanni) Adewoyin <[email protected]> * test: adjust to dropout_p * doc: update changelog * Feature/integrate reusable workflows (#16) * ci: add public pr label * ci: add readthedocs update check * ci: add downstream ci * ci: add ci-config * chore(deps): remove unused dependency * docs: update changelog * ci: switch to main * chore: changelog 0.2.1 * Update error messages from invalid sub_graph in model instantiation (#20) * ci: inherit pypi publish flow (#17) * ci: inherit pypi publish flow Co-authored-by: Helen Theissen <[email protected]> * docs: add to changelog * fix: typo in reusable workflow * fix: another typo * chore: bump actions/setup-python to v5 * ci: run downstream-ci for changes in src and tests * docs: update changelog --------- Co-authored-by: Helen Theissen <[email protected]> * Update CHANGELOG.md to KeepChangelog format * [pre-commit.ci] pre-commit autoupdate (#25) updates: - [github.com/psf/black-pre-commit-mirror: 24.4.2 → 24.8.0](psf/black-pre-commit-mirror@24.4.2...24.8.0) - [github.com/astral-sh/ruff-pre-commit: v0.4.6 → v0.6.2](astral-sh/ruff-pre-commit@v0.4.6...v0.6.2) - [github.com/tox-dev/pyproject-fmt: 2.1.3 → 2.2.1](tox-dev/pyproject-fmt@2.1.3...2.2.1) Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> * Ci/changelog-release-updater (#26) * ci: add changelof release updater * docs: update changelog * Feature/integrate reusable workflows (#16) * ci: add public pr label * ci: add readthedocs update check * ci: add downstream ci * ci: add ci-config * chore(deps): remove unused dependency * docs: update changelog * ci: switch to main * chore: changelog 0.2.1 * Update error messages from invalid sub_graph in model instantiation (#20) * ci: inherit pypi publish flow (#17) * ci: inherit pypi publish flow Co-authored-by: Helen Theissen <[email protected]> * docs: add to changelog * fix: typo in reusable workflow * fix: another typo * chore: bump actions/setup-python to v5 * ci: run downstream-ci for changes in src and tests * docs: update changelog --------- Co-authored-by: Helen Theissen <[email protected]> * Update CHANGELOG.md to KeepChangelog format * Ci/changelog-release-updater (#26) * ci: add changelof release updater * docs: update changelog --------- Co-authored-by: Rilwan (Akanni) Adewoyin <[email protected]> Co-authored-by: Gert Mertes <[email protected]> Co-authored-by: Mario Santa Cruz <[email protected]> Co-authored-by: Jesper Dramsch <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

xfail for MultiHeadSelfAttention

for more information, see https://pre-commit.ci

codecov-commenter · 2024-10-01T08:14:35Z

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 99.84%. Comparing base (e608a73) to head (fc335c7).
Report is 12 commits behind head on develop.

Additional details and impacted files

@@           Coverage Diff            @@
##           develop      #47   +/-   ##
========================================
  Coverage    99.84%   99.84%           
========================================
  Files           23       23           
  Lines         1277     1278    +1     
========================================
+ Hits          1275     1276    +1     
  Misses           2        2

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

CHANGELOG.md

mchantry · 2024-10-03T10:51:08Z

Where's the PR template?

src/anemoi/models/layers/attention.py

theissenhelen linked an issue Sep 20, 2024 that may be closed by this pull request

Make Flash attention configurable #44

Open

mchantry reviewed Sep 22, 2024

View reviewed changes

src/anemoi/models/layers/attention.py Outdated Show resolved Hide resolved

theissenhelen and others added 9 commits September 27, 2024 08:42

feat: FlashMultiHeadSelfAttention

539e8a2

chore!: drop support for scaled_dot_product_attention

a86c9a8

feat: add softcap

105443f

test: add softcap

e82a59e

xfail for MultiHeadSelfAttention

[pre-commit.ci] auto fixes from pre-commit.com hooks

e648eb0

for more information, see https://pre-commit.ci

feat: flash attention lazy import

6271cd8

feat: make alibi slopes configurable

d4940e7

theissenhelen force-pushed the feature/44-make-flash-attention-configurable branch from a080cc5 to d4940e7 Compare September 27, 2024 08:44

theissenhelen added 4 commits September 27, 2024 11:49

chore(deps): add flash-attn

9ff6cb9

feat: use scaled_dot_product as default

bbd89dc

feat: make alibi_slope cinfigurable in block, chunk processor

91533c6

chore(deps): remove flash-attn

0eb5c50

theissenhelen added 6 commits October 2, 2024 13:07

feat: get alibi_slopes

c04e641

docs: update docstrings

6523b47

fix: bias shape

22623cc

fix: softcap optional

ed07e34

fix: import annotations from future

c841324

fix: annotation error

6c12dda

theissenhelen marked this pull request as ready for review October 3, 2024 10:08

theissenhelen requested a review from mchantry October 3, 2024 10:08

mchantry reviewed Oct 3, 2024

View reviewed changes

CHANGELOG.md Show resolved Hide resolved

docs: update changelog

b7b8f2e

theissenhelen mentioned this pull request Oct 3, 2024

feat: make flash-attn configurabele ecmwf/anemoi-training#73

Draft

mpvginde reviewed Oct 4, 2024

View reviewed changes

src/anemoi/models/layers/attention.py Outdated Show resolved Hide resolved

theissenhelen added 2 commits October 7, 2024 07:33

fix: type annotation

df353d9

feat: catch low flash-attn version

fc335c7

theissenhelen requested review from mchantry and mpvginde October 7, 2024 11:58

mpvginde approved these changes Oct 7, 2024

View reviewed changes

b8raoult reviewed Oct 15, 2024

View reviewed changes

src/anemoi/models/layers/attention.py Show resolved Hide resolved

theissenhelen requested a review from b8raoult October 21, 2024 15:37

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Feature/44 make flash attention configurable #47

Feature/44 make flash attention configurable #47

theissenhelen commented Sep 20, 2024 •

edited

Loading

codecov-commenter commented Oct 1, 2024 •

edited

Loading

mchantry commented Oct 3, 2024

Feature/44 make flash attention configurable #47

Are you sure you want to change the base?

Feature/44 make flash attention configurable #47

Conversation

theissenhelen commented Sep 20, 2024 • edited Loading

codecov-commenter commented Oct 1, 2024 • edited Loading

Codecov Report

mchantry commented Oct 3, 2024

theissenhelen commented Sep 20, 2024 •

edited

Loading

codecov-commenter commented Oct 1, 2024 •

edited

Loading