Feature/44 make flash attention configurable #60

theissenhelen · 2025-01-06T16:27:11Z

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:

flex attention available
user specifies whether flash-attn, flex attention or scaled dot product attention 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
scaled dot product attention now supports sliding window (making it numerically equivalent to flash/flex)

Todo:

test various attention options
adjust test case coverage

* 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

* 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

….com:ecmwf/anemoi-models into feature/44-make-flash-attention-configurable

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

….com:ecmwf/anemoi-models into feature/44-make-flash-attention-configurable

…44-make-flash-attention-configurable

FussyDuck · 2025-01-06T16:27:17Z

All committers have signed the CLA.

theissenhelen and others added 30 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

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

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

docs: update changelog

b7b8f2e

fix: type annotation

df353d9

feat: catch low flash-attn version

fc335c7

feat: FlashMultiHeadSelfAttention

663fea0

chore!: drop support for scaled_dot_product_attention

0c55a9c

feat: add softcap

ea665be

test: add softcap

ffa2d99

xfail for MultiHeadSelfAttention

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

7c2d634

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

feat: flash attention lazy import

d2ed932

cathalobrien and others added 27 commits November 1, 2024 14:42

added flex attn wrapper

147e772

fix: alibi_slopes unassigned

f0c24e8

adding causal wip

3c4572b

Merge branch 'feature/44-make-flash-attention-configurable' of github…

fb731f7

….com:ecmwf/anemoi-models into feature/44-make-flash-attention-configurable

added flex attn module

f0308f2

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

6dee265

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

Bump min torch version to be able to use Flex Attn

7fb0b62

added input parameter checks

739aa65

precommit fix

2a2ed11

merge

fa1474c

fix: typo

a703688

test: adjust tests

f1be563

fix: no self.use_alibi_slopes

0dda5d6

fix: use_alibi_slope default to false

12facf0

feat: Add sliding window support for TorchAttention via mask

60e32f1

fix: set default flash_attention

07d9684

fix: pytest

9a1827a

fix: tests

ca8c9fa

Merge branch 'feature/44-make-flash-attention-configurable' of github…

ac897ea

….com:ecmwf/anemoi-models into feature/44-make-flash-attention-configurable

docs: improve docstrings in MultiHeadSelfAttention

7ec8142

fix: error instead of SystemExit

972d3c5

chore: refactor SDPAAttention update_mask method

e89fd2e

feat: add missing pytest.ini

2d122df

chore: remove explicit float typing

d4510f6

Merge branch 'feature/44-make-flash-attention-configurable' of github…

6057004

….com:ecmwf/anemoi-models into feature/44-make-flash-attention-configurable

support running without window size

8656cae

Merge commit '8656cae768f8f8359772cdab0b06e0933b8a0bdf' into feature/…

705fc6b

…44-make-flash-attention-configurable

anaprietonem added the models label Jan 9, 2025

anaprietonem assigned theissenhelen Jan 9, 2025

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Feature/44 make flash attention configurable #60

Feature/44 make flash attention configurable #60

theissenhelen commented Jan 6, 2025

FussyDuck commented Jan 6, 2025 •

edited

Loading

Feature/44 make flash attention configurable #60

Are you sure you want to change the base?

Feature/44 make flash attention configurable #60

Conversation

theissenhelen commented Jan 6, 2025

FussyDuck commented Jan 6, 2025 • edited Loading

FussyDuck commented Jan 6, 2025 •

edited

Loading