Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

add flash-attn deterministic option to flash-attn>=2.4.1 #31961

Merged
merged 4 commits into from
Jul 16, 2024
Merged

add flash-attn deterministic option to flash-attn>=2.4.1 #31961

merged 4 commits into from
Jul 16, 2024

Conversation

junrae6454
Copy link
Contributor

What does this PR do?

This PR introduces support for the deterministic mode available in Flash Attention version 2.4.1 and above.
The deterministic mode in Flash Attention 2.4.1 ensures reproducible results, which is crucial for debugging and scientific research. By integrating this feature, users can benefit from deterministic behavior in their models.
The changes include enabling deterministic behavior through an environment variable and integrating this feature into the _flash_attention_forward function.

  • Flash Attention Forward Function Update:

    • Added deterministic parameter to _flash_attention_forward to control deterministic behavior.
    • Modified the function to check for Flash Attention version 2.4.1 or higher and set the deterministic argument accordingly.

  • Trainer Utilities Update:

    • Set the FLASH_ATTENTION_DETERMINISTIC environment variable in enable_full_determinism.

Before submitting

  • This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
  • Did you read the contributor guideline,
    Pull Request section?
  • Was this discussed/approved via a Github issue or the forum? Please add a link
    to it if that's the case.
  • Did you make sure to update the documentation with your changes? Here are the
    documentation guidelines, and
    here are tips on formatting docstrings.
  • Did you write any new necessary tests?

Who can review?

@ArthurZucker

ArthurZucker
ArthurZucker reviewed Jul 15, 2024
View reviewed changes
Copy link
Collaborator

@ArthurZucker ArthurZucker left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

sounds good! are there other args we should be reading from env ?

src/transformers/modeling_flash_attention_utils.py Outdated
from typing import Optional, Tuple

import torch
import torch.nn.functional as F

from .utils import is_flash_attn_2_available
from .utils import is_flash_attn_2_available, is_flash_attn_greater_or_equal_2_41
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
from .utils import is_flash_attn_2_available, is_flash_attn_greater_or_equal_2_41
from .utils import is_flash_attn_2_available, is_flash_attn_greater

src/transformers/modeling_flash_attention_utils.py Outdated
@@ -177,6 +181,9 @@ def _flash_attention_forward(
)
flash_kwargs = {"window_size": (sliding_window, sliding_window)} if use_sliding_windows else {}

if is_flash_attn_greater_or_equal_2_41():
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
if is_flash_attn_greater_or_equal_2_41():
if is_flash_attn_greater_or_equal("2.41"):

src/transformers/utils/import_utils.py Outdated
Comment on lines 817 to 821
def is_flash_attn_greater_or_equal_2_41():
if not _is_package_available("flash_attn"):
return False

return version.parse(importlib.metadata.version("flash_attn")) >= version.parse("2.4.1")
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
def is_flash_attn_greater_or_equal_2_41():
if not _is_package_available("flash_attn"):
return False
return version.parse(importlib.metadata.version("flash_attn")) >= version.parse("2.4.1")

unnecessary, see def is_flash_attn_greater_or_equal(library_version: str):

@junrae6454
Copy link
Contributor Author

sounds good! are there other args we should be reading from env ?

@ArthurZucker
Other changes are related to the model structure and the model's return value, so there doesn't seem to be anything else that needs to be read from the environment.

@ArthurZucker ArthurZucker mentioned this pull request Jul 16, 2024
Closed
4 tasks
ArthurZucker
ArthurZucker approved these changes Jul 16, 2024
View reviewed changes
Copy link
Collaborator

@ArthurZucker ArthurZucker left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

thanks for iterating

@ArthurZucker ArthurZucker merged commit 036d3de into huggingface:main Jul 16, 2024
18 of 20 checks passed
@Cyrilvallez Cyrilvallez mentioned this pull request Jul 17, 2024
Merged
amyeroberts pushed a commit to amyeroberts/transformers that referenced this pull request Jul 19, 2024
@junrae6454 @amyeroberts
add flash-attn deterministic option to flash-attn>=2.4.1 (huggingface…
abb40e4 
…#31961)

* add flash-attn deterministic option to flash-attn>=2.4.1

* Add Missing Import

* Fix ruff linting issues

* Replace `is_flash_attn_greater_or_equal_2_41` with the existing `is_flash_attn_greater_or_equal`

---------

Co-authored-by: jun.4 <[email protected]>
@junrae6454 junrae6454 mentioned this pull request Jul 23, 2024
Merged
5 tasks
MHRDYN7 pushed a commit to MHRDYN7/transformers that referenced this pull request Jul 23, 2024
@junrae6454 @MHRDYN7
add flash-attn deterministic option to flash-attn>=2.4.1 (huggingface…
00e0b24 
…#31961)

* add flash-attn deterministic option to flash-attn>=2.4.1

* Add Missing Import

* Fix ruff linting issues

* Replace `is_flash_attn_greater_or_equal_2_41` with the existing `is_flash_attn_greater_or_equal`

---------

Co-authored-by: jun.4 <[email protected]>
zucchini-nlp pushed a commit to zucchini-nlp/transformers that referenced this pull request Jul 24, 2024
@junrae6454 @zucchini-nlp
add flash-attn deterministic option to flash-attn>=2.4.1 (huggingface…
722ac3a 
…#31961)

* add flash-attn deterministic option to flash-attn>=2.4.1

* Add Missing Import

* Fix ruff linting issues

* Replace `is_flash_attn_greater_or_equal_2_41` with the existing `is_flash_attn_greater_or_equal`

---------

Co-authored-by: jun.4 <[email protected]>
itazap pushed a commit that referenced this pull request Jul 25, 2024
@junrae6454 @itazap
add flash-attn deterministic option to flash-attn>=2.4.1 (#31961)
ad56f03 
* add flash-attn deterministic option to flash-attn>=2.4.1

* Add Missing Import

* Fix ruff linting issues

* Replace `is_flash_attn_greater_or_equal_2_41` with the existing `is_flash_attn_greater_or_equal`

---------

Co-authored-by: jun.4 <[email protected]>
@ArthurZucker ArthurZucker mentioned this pull request Aug 5, 2024
Closed
4 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@Boubou78000 Boubou78000 Boubou78000 approved these changes

@ArthurZucker ArthurZucker ArthurZucker approved these changes

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
3 participants
@junrae6454 @ArthurZucker @Boubou78000