[ROCM] Properly disable Flash Attention/Efficient Attention with environment variables by xinyazhang · Pull Request #133866 · pytorch/pytorch

xinyazhang · 2024-08-19T16:38:39Z

Now USE_FLASH_ATTENTION=0 USE_MEM_EFF_ATTENTION=0 python setup.py can compile correctly

Fixes #125230

cc @jeffdaily @sunway513 @jithunnair-amd @pruthvistony @ROCmSupport @dllehr-amd @jataylo @hongxiayang

…NTION=0

…M_EFF_ATTENTION=0

pytorch-bot · 2024-08-19T16:38:42Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/133866

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (3 Unrelated Failures)

As of commit 298fee6 with merge base df68315 ():

FLAKY - The following job failed but was likely due to flakiness present on trunk:

rocm / linux-focal-rocm6.1-py3.8 / test (default, 2, 6, linux.rocm.gpu.2) (gh) (similar failure)
'test/test_transformers.py::TestSDPACudaOnlyCUDA::test_flash_attention_vs_math_ref_grads_batch_size_1_seq_len_q_143_seq_len_k_127_head_dim_203_is_causal_False_dropout_p_0_0_bfloat16_scale0_enable_gqa_True_n_heads0_cuda_bfloat16'

BROKEN TRUNK - The following jobs failed but was present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

This comment was automatically generated by Dr. CI and updates every 15 minutes.

jithunnair-amd

LGTM. @xinyazhang Have you done any local testing with a build that has them disabled?

xinyazhang · 2024-08-19T20:10:17Z

LGTM. @xinyazhang Have you done any local testing with a build that has them disabled?

I run a local sdpa.py benchmark script (got this from slack and unsuitable to show it here) and force it to use FA. The script aborted as expected after reporting FA not supported.

jeffdaily · 2024-08-21T19:33:04Z

CI build failure is real.

2024-08-20T06:28:34.5898661Z In file included from /var/lib/jenkins/workspace/aten/src/ATen/native/transformers/hip/attention_backward.hip:49:
2024-08-20T06:28:34.5900419Z /var/lib/jenkins/workspace/aten/src/ATen/native/transformers/hip/aotriton_adapter.h:5:10: fatal error: 'aotriton/dtypes.h' file not found
2024-08-20T06:28:34.5901514Z     5 | #include <aotriton/dtypes.h>
2024-08-20T06:28:34.5901901Z       |          ^~~~~~~~~~~~~~~~~~~
2024-08-20T06:28:34.5902317Z 1 error generated when compiling for host.

xinyazhang · 2024-08-21T21:41:19Z

CI build failure is real.

@jeffdaily It is expected if AOTriton installation doesn't exist (which I think it's the case for the CI image)

2024-08-20T06:20:33.9782632Z --     ROCM_VERSION        : 
2024-08-20T06:20:33.9783026Z --     USE_FLASH_ATTENTION : ON
2024-08-20T06:20:33.9783425Z --     USE_MEM_EFF_ATTENTION : ON

Both should be OFF to disable AOTriton.

xinyazhang · 2024-08-21T23:07:41Z

Okay I think I found the problem. Within Dependencies.cmake, neither USE_FLASH_ATTENTION nor USE_MEM_EFF_ATTENTION is defined.

Otherwise we have circular dependencies.

pruthvistony · 2024-08-22T22:32:08Z

@malfet ,
Please help on review of this PR.

CMakeLists.txt

jithunnair-amd · 2024-08-22T22:53:23Z

@pytorchbot merge

pytorchmergebot · 2024-08-22T22:55:01Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

pytorchmergebot · 2024-08-23T04:53:51Z

The merge job was canceled or timed out. This most often happen if two merge requests were issued for the same PR, or if merge job was waiting for more than 6 hours for tests to finish. In later case, please do not hesitate to reissue the merge command
For more information see pytorch-bot wiki.

xinyazhang · 2024-08-23T15:06:54Z

@jithunnair-amd it seems the merge was blocked by failing CI, which is supposed be fixed by #133884

jithunnair-amd · 2024-08-27T18:22:47Z

@pytorchbot merge -f "Build issues resolved. This PR is for build scenarios not relevant to CI. Test failures are related to GQA which is addressed in #133884."

pytorchmergebot · 2024-08-27T18:24:17Z

Merge started

Your change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes). Please use -f as last resort and instead consider -i/--ignore-current to continue the merge ignoring current failures. This will allow currently pending tests to finish and report signal before the merge.

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

…tion with environment variables (#1570) Now `USE_FLASH_ATTENTION=0 USE_MEM_EFF_ATTENTION=0 python setup.py` can compile correctly. This is cherry-picked version of pytorch#133866

…tion with environment variables (#1571) Now `USE_FLASH_ATTENTION=0 USE_MEM_EFF_ATTENTION=0 python setup.py` can compile correctly. This is cherry-picked version of pytorch#133866 --------- Co-authored-by: Pruthvi Madugundu <pruthvigithub@gmail.com>

…ronment variables (#1542) Now `USE_FLASH_ATTENTION=0 USE_MEM_EFF_ATTENTION=0 python setup.py` can compile correctly. This is cherry-picked version of pytorch#133866 Tested with `USE_FLASH_ATTENTION=0 USE_MEM_EFF_ATTENTION=0 python setup.py develop --user` and `python -c 'import torch'`

…ronment variables (pytorch#133866) Now `USE_FLASH_ATTENTION=0 USE_MEM_EFF_ATTENTION=0 python setup.py` can compile correctly Fixes pytorch#125230 Pull Request resolved: pytorch#133866 Approved by: https://github.com/jithunnair-amd, https://github.com/jeffdaily, https://github.com/malfet

…tion with environment variables (#1570) Now `USE_FLASH_ATTENTION=0 USE_MEM_EFF_ATTENTION=0 python setup.py` can compile correctly. This is cherry-picked version of pytorch#133866

xinyazhang added 2 commits August 19, 2024 16:04

Do not include aotriton if USE_FLASH_ATTENTION=0 AND USE_MEM_EFF_ATTE…

a7c8ef6

…NTION=0

Disable AOTriton in sdp_utils.cpp if USE_FLASH_ATTENTION=0 AND USE_ME…

50d2588

…M_EFF_ATTENTION=0

pytorch-bot bot added the module: rocm AMD GPU support for Pytorch label Aug 19, 2024

pytorchbot added the open source label Aug 19, 2024

This was referenced Aug 19, 2024

[ROCM] Properly disable Flash Attention/Efficient Attention with environment variables ROCm/pytorch#1541

Closed

[ROCM] Properly disable Flash Attention/Efficient Attention with environment variables ROCm/pytorch#1542

Merged

jithunnair-amd approved these changes Aug 19, 2024

View reviewed changes

jithunnair-amd approved these changes Aug 21, 2024 •

edited

Loading

View reviewed changes

xinyazhang marked this pull request as ready for review August 21, 2024 15:02

jeffdaily approved these changes Aug 21, 2024

View reviewed changes

jithunnair-amd mentioned this pull request Aug 21, 2024

[ROCm] Allow to disable USE_FLASH_ATTENTION/USE_MEM_EFF_ATTENTION #130197

Closed

Move include(cmake/External/aotriton.cmake) out of Dependencies.cmake

693d4e3

Otherwise we have circular dependencies.

xinyazhang force-pushed the xinyazhang/nofa-2.5main branch from e767d4b to 693d4e3 Compare August 21, 2024 23:08

Remove the legacy code that sets USE_MEM_EFF_ATTENTION to OFF

298fee6

pruthvistony added the rocm This tag is for PRs from ROCm team label Aug 22, 2024

jithunnair-amd requested a review from malfet August 22, 2024 22:32

malfet added the ciflow/rocm Trigger "default" config CI on ROCm label Aug 22, 2024

malfet approved these changes Aug 22, 2024

View reviewed changes

malfet reviewed Aug 22, 2024

View reviewed changes

CMakeLists.txt Show resolved Hide resolved

CMakeLists.txt Show resolved Hide resolved

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Aug 22, 2024

jithunnair-amd added the topic: not user facing topic category label Aug 22, 2024

pytorchmergebot added the merging label Aug 22, 2024

pytorchmergebot added the Merged label Aug 27, 2024

pytorchmergebot closed this in 5fd670e Aug 27, 2024

pytorchmergebot removed the merging label Aug 27, 2024

This was referenced Sep 3, 2024

[release/2.4] [ROCM] Properly disable Flash Attention/Efficient Attention with environment variables ROCm/pytorch#1570

Merged

[release/2.3] [ROCM] Properly disable Flash Attention/Efficient Attention with environment variables ROCm/pytorch#1571

Merged

Conversation

xinyazhang commented Aug 19, 2024 • edited by pytorch-bot bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Aug 19, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/133866

✅ You can merge normally! (3 Unrelated Failures)

Uh oh!

jithunnair-amd left a comment

Choose a reason for hiding this comment

Uh oh!

xinyazhang commented Aug 19, 2024

Uh oh!

Uh oh!

jeffdaily commented Aug 21, 2024

Uh oh!

xinyazhang commented Aug 21, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

xinyazhang commented Aug 21, 2024

Uh oh!

pruthvistony commented Aug 22, 2024

Uh oh!

Uh oh!

Uh oh!

jithunnair-amd commented Aug 22, 2024

Uh oh!

pytorchmergebot commented Aug 22, 2024

Merge started

Uh oh!

pytorchmergebot commented Aug 23, 2024

Uh oh!

xinyazhang commented Aug 23, 2024

Uh oh!

jithunnair-amd commented Aug 27, 2024

Uh oh!

pytorchmergebot commented Aug 27, 2024

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

xinyazhang commented Aug 19, 2024 •

edited by pytorch-bot bot

Loading

pytorch-bot bot commented Aug 19, 2024 •

edited

Loading

xinyazhang commented Aug 21, 2024 •

edited

Loading