Skip to content
This repository has been archived by the owner on Dec 20, 2024. It is now read-only.

Feature/44 make flash attention configurable #47

Open
wants to merge 73 commits into
base: develop
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
73 commits
Select commit Hold shift + click to select a range
539e8a2
feat: FlashMultiHeadSelfAttention
theissenhelen Sep 17, 2024
3317138
Chore/multiple fixes ci precommit (#41)
theissenhelen Sep 18, 2024
3186a8e
11 add configurability to dropout in multiheadselfattention module (#12)
theissenhelen Sep 18, 2024
a86c9a8
chore!: drop support for scaled_dot_product_attention
theissenhelen Sep 20, 2024
105443f
feat: add softcap
theissenhelen Sep 20, 2024
e82a59e
test: add softcap
theissenhelen Sep 20, 2024
e648eb0
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Sep 20, 2024
6271cd8
feat: flash attention lazy import
theissenhelen Sep 23, 2024
d4940e7
feat: make alibi slopes configurable
theissenhelen Sep 27, 2024
9ff6cb9
chore(deps): add flash-attn
theissenhelen Sep 27, 2024
bbd89dc
feat: use scaled_dot_product as default
theissenhelen Oct 1, 2024
91533c6
feat: make alibi_slope cinfigurable in block, chunk processor
theissenhelen Oct 1, 2024
0eb5c50
chore(deps): remove flash-attn
theissenhelen Oct 1, 2024
c04e641
feat: get alibi_slopes
theissenhelen Oct 2, 2024
6523b47
docs: update docstrings
theissenhelen Oct 3, 2024
22623cc
fix: bias shape
theissenhelen Oct 3, 2024
ed07e34
fix: softcap optional
theissenhelen Oct 3, 2024
c841324
fix: import annotations from future
theissenhelen Oct 3, 2024
6c12dda
fix: annotation error
theissenhelen Oct 3, 2024
b7b8f2e
docs: update changelog
theissenhelen Oct 3, 2024
df353d9
fix: type annotation
theissenhelen Oct 7, 2024
fc335c7
feat: catch low flash-attn version
theissenhelen Oct 7, 2024
663fea0
feat: FlashMultiHeadSelfAttention
theissenhelen Sep 17, 2024
a8b3f9d
Chore/multiple fixes ci precommit (#41)
theissenhelen Sep 18, 2024
6595ca1
11 add configurability to dropout in multiheadselfattention module (#12)
theissenhelen Sep 18, 2024
0c55a9c
chore!: drop support for scaled_dot_product_attention
theissenhelen Sep 20, 2024
ea665be
feat: add softcap
theissenhelen Sep 20, 2024
ffa2d99
test: add softcap
theissenhelen Sep 20, 2024
7c2d634
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Sep 20, 2024
d2ed932
feat: flash attention lazy import
theissenhelen Sep 23, 2024
3295159
feat: make alibi slopes configurable
theissenhelen Sep 27, 2024
ebde686
chore(deps): add flash-attn
theissenhelen Sep 27, 2024
5102d9a
feat: use scaled_dot_product as default
theissenhelen Oct 1, 2024
3abc286
feat: make alibi_slope cinfigurable in block, chunk processor
theissenhelen Oct 1, 2024
673a25d
chore(deps): remove flash-attn
theissenhelen Oct 1, 2024
f606058
feat: get alibi_slopes
theissenhelen Oct 2, 2024
ef34771
docs: update docstrings
theissenhelen Oct 3, 2024
5136fb3
fix: bias shape
theissenhelen Oct 3, 2024
892c269
fix: softcap optional
theissenhelen Oct 3, 2024
4c42171
fix: import annotations from future
theissenhelen Oct 3, 2024
4bdf464
fix: annotation error
theissenhelen Oct 3, 2024
5a670b2
docs: update changelog
theissenhelen Oct 3, 2024
34db6e4
fix: type annotation
theissenhelen Oct 7, 2024
d424c75
feat: catch low flash-attn version
theissenhelen Oct 7, 2024
222b7d8
feat: attention wrapper
theissenhelen Oct 25, 2024
c2aca14
fix: remove duplicate version check
theissenhelen Oct 25, 2024
b75d225
merge conflict
cathalobrien Nov 1, 2024
147e772
added flex attn wrapper
cathalobrien Nov 1, 2024
f0c24e8
fix: alibi_slopes unassigned
theissenhelen Nov 6, 2024
3c4572b
adding causal wip
cathalobrien Nov 6, 2024
fb731f7
Merge branch 'feature/44-make-flash-attention-configurable' of github…
cathalobrien Nov 8, 2024
f0308f2
added flex attn module
cathalobrien Nov 12, 2024
6dee265
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Nov 12, 2024
7fb0b62
Bump min torch version to be able to use Flex Attn
cathalobrien Nov 12, 2024
739aa65
added input parameter checks
cathalobrien Nov 12, 2024
2a2ed11
precommit fix
cathalobrien Nov 12, 2024
fa1474c
merge
cathalobrien Nov 12, 2024
a703688
fix: typo
theissenhelen Nov 26, 2024
f1be563
test: adjust tests
theissenhelen Nov 27, 2024
0dda5d6
fix: no self.use_alibi_slopes
theissenhelen Nov 27, 2024
12facf0
fix: use_alibi_slope default to false
theissenhelen Nov 28, 2024
60e32f1
feat: Add sliding window support for TorchAttention via mask
japols Dec 9, 2024
07d9684
fix: set default flash_attention
japols Dec 10, 2024
9a1827a
fix: pytest
japols Dec 10, 2024
ca8c9fa
fix: tests
japols Dec 13, 2024
ac897ea
Merge branch 'feature/44-make-flash-attention-configurable' of github…
cathalobrien Dec 16, 2024
7ec8142
docs: improve docstrings in MultiHeadSelfAttention
theissenhelen Dec 18, 2024
972d3c5
fix: error instead of SystemExit
theissenhelen Dec 18, 2024
e89fd2e
chore: refactor SDPAAttention update_mask method
theissenhelen Dec 18, 2024
2d122df
feat: add missing pytest.ini
theissenhelen Dec 18, 2024
d4510f6
chore: remove explicit float typing
theissenhelen Dec 19, 2024
6057004
Merge branch 'feature/44-make-flash-attention-configurable' of github…
cathalobrien Dec 19, 2024
8656cae
support running without window size
cathalobrien Dec 19, 2024
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -39,15 +39,24 @@ Keep it human-readable, your future self will thank you!

### Added

- CI workflow to update the changelog on release
theissenhelen marked this conversation as resolved.
Show resolved Hide resolved
- add configurability of flash attention (#47)
- configurabilty of the dropout probability in the the MultiHeadSelfAttention module
- CI workflow to update the changelog on release
- Remapper: Preprocessor for remapping one variable to multiple ones. Includes changes to the data indices since the remapper changes the number of variables. With optional config keywords.
- Codeowners file
- Pygrep precommit hooks
- Docsig precommit hooks
- Changelog merge strategy

theissenhelen marked this conversation as resolved.
Show resolved Hide resolved

### Changed

- Update CI to inherit from common infrastructue reusable workflows
- run downstream-ci only when src and tests folders have changed
- New error messages for wrongs graphs.
- Feature: Change model to be instantiatable in the interface, addressing [#28](https://github.com/ecmwf/anemoi-models/issues/28) through [#45](https://github.com/ecmwf/anemoi-models/pulls/45)
- Bugfixes for CI

### Removed

Expand Down
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@ dependencies = [
"anemoi-utils>=0.1.9",
"einops>=0.6.1",
"hydra-core>=1.3",
"torch>=2.2",
"torch>=2.5",
"torch-geometric>=2.3,<2.5",
]
optional-dependencies.all = [ ]
Expand Down
7 changes: 7 additions & 0 deletions pytest.ini
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
[pytest]
markers =
data_dependent: marks tests depending on data (deselect with '-m "not data_dependent"')
auth: marks tests that require authentication (deselect with '-m "not auth"')
gpu: marks tests that require a GPU (deselect with '-m "not gpu"')

tmp_path_retention_policy = none
Loading
Loading