Skip to content

[Backport] Add flash attention as custom op so dynamo can trace it (#6875)#6884

Merged
JackCaoG merged 4 commits intor2.3from
JackCaoG/backport_flash_attention_dynamo
Apr 5, 2024
Merged

[Backport] Add flash attention as custom op so dynamo can trace it (#6875)#6884
JackCaoG merged 4 commits intor2.3from
JackCaoG/backport_flash_attention_dynamo

Conversation

@JackCaoG
Copy link
Copy Markdown
Collaborator

@JackCaoG JackCaoG commented Apr 4, 2024

We intended to add dynamo support for pallas but flash attention support needs to be added.

@JackCaoG JackCaoG requested review from alanwaketan and lsy323 April 4, 2024 00:02
alanwaketan
alanwaketan approved these changes Apr 4, 2024
View reviewed changes
Copy link
Copy Markdown
Collaborator

@alanwaketan alanwaketan left a comment

Choose a reason for hiding this comment

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

LGTM.

@JackCaoG JackCaoG mentioned this pull request Apr 4, 2024
Closed
@JackCaoG
Copy link
Copy Markdown
Collaborator Author

JackCaoG commented Apr 4, 2024

Ehh TPU test failed with

======================================================================
FAIL: test_resnet18 (__main__.DynamoTrainingOptimizerTest)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/home/runner/_work/xla/xla/pytorch/xla/test/dynamo/test_dynamo.py", line 519, in test_resnet18
    self.assertEqual(met.metric_data('CompileTime')[0], 7)
AssertionError: 5 != 7

@lsy323 @will-cromar I think TPU CI is incorrectly always checkout the PyTorch's main branch and ignore the .torch_pin

@will-cromar
Copy link
Copy Markdown
Collaborator

will-cromar commented Apr 4, 2024

Can you just hardcode the release branch on this branch?

xla/.github/workflows/tpu_ci.yml

Line 24 in 7520c04

git clone --recursive https://github.com/pytorch/pytorch

I can update master to use the ansible config so we don't have to do stuff like this twice.

@JackCaoG
Copy link
Copy Markdown
Collaborator Author

JackCaoG commented Apr 5, 2024

I am going to merge this one as this also fix the TPUCI branch for the release branch.

@JackCaoG JackCaoG merged commit 96cc98f into r2.3 Apr 5, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@alanwaketan alanwaketan alanwaketan approved these changes

@lsy323 lsy323 Awaiting requested review from lsy323

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

@JackCaoG @will-cromar @alanwaketan