Add FMHA T5x test #442
Merged
Add FMHA T5x test #442
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Closed
|
Closed the old PR #327 as this one is the latest one. |
terrykong
requested changes
Jan 28, 2024
|
Moved to draft since there still seems to be active development on this PR. BTW, it may be better to add this to _test_t5x.yaml since we were already running pretty close to the maximum number of reusable workflows allowed in github. |
hmonishN
force-pushed
the
add_t5x_fmha_test
branch
from
abdf793
to
a5c2f5d
Compare
hmonishN
force-pushed
the
add_t5x_fmha_test
branch
from
e8d968a
to
89a47cb
Compare
Merginf fmha into t5x test.
Merginf fmha into t5x.
hmonishN
force-pushed
the
add_t5x_fmha_test
branch
from
89a47cb
to
cc95535
Compare
terrykong
reviewed
Mar 7, 2024
terrykong
requested changes
Mar 7, 2024
Incorporated review comments
fixed indentation
terrykong
reviewed
Mar 7, 2024
updating the new name for _test_t5x Co-authored-by: Terry Kong <terryk@nvidia.com>
reduced number of jobs for fmha
terrykong
reviewed
Mar 7, 2024
terrykong
approved these changes
Mar 8, 2024
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Adding the JAX T5x FMHA E2E system test to check for fmha lowering support. Following are the steps implemented in the test:
FMHA lowering flag is enabled by default now, enabled the dumping of hlo to track fmha forward and backward instructions.
Added the test as part of _ci.yaml file and also added a nightly workflow file for it. We will add this test as part of performance benchmarking later and add hlo to baseline.
Also added changes for correction of seq length of decoder (should be a multiple of 64)
The test was failing with following error related to CUDNN_STATUS_BAD_PARAM. The fix for this is added in the [PR] (openxla/xla#6872) in upstream which is now merged and the test passes. Bug for this error.
run for these changes: workflow run link