Add flash_attn_3 as first option for FA3 import - #2010
Conversation
Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
|
I was shocked that FA3 is not in the NGC pytorch containers. Is it installed under this alternate names? |
| try: | ||
| from flashattn_hopper.flash_attn_interface import _flash_attn_forward | ||
| from flashattn_hopper.flash_attn_interface import ( | ||
| from flash_attn_3.flash_attn_interface import _flash_attn_forward |
There was a problem hiding this comment.
Can we get rid of the nested try/except clause loops?
It would look a bit nicer if we have smth like:
HAVE_FA3=False
try:
# import old way
HAVE_FA3=True
except:
pass # ideally log
try:
# import new way
HAVE_FA3=True
except:
pass # ideally logOr maybe add a branch depending on the fa version or smth. TransformerEngine has FlashAttentionUtils that can give you the FA version, we could use that to unify the behaviour between the training and inference side.
There was a problem hiding this comment.
The second try block would have to be wrapped in a if not HAVE_FA3 clause though, so it seems like it would be equally messy? Open to either option though
There was a problem hiding this comment.
Agreed! Still looks more linear/simpler to me.
There was a problem hiding this comment.
Updated to not nest the imports
|
/ok to test a4410fc |
|
/ok to test 3eae66f |
7 participants
What does this PR do ?
Makes
flash_attn_3the first option for importing FlashAttention 3 as this is the default location.Contribution process
flowchart LR A[Pre-checks] --> B[PR Tests] subgraph Code Review/Approval C1[Expert Review] --> C2[Final Review] end B --> C1 C2 --> D[Merge]Pre-checks
Core 0.8)Code review
The following process is enforced via the CODEOWNERS file for changes into
megatron/core. For changes outside ofmegatron/core, it is up to the PR author whether or not to tag the Final Reviewer team.For MRs into `main` branch
(Step 1): Add PR label
Expert Review(Step 2): Collect the expert reviewers reviews
Expert Reviewlabel when your PR is ready for review.Final Review might get declined if these requirements are not fulfilled.
(Step 3): Final Review
Final Reviewlabel(Optional Step 4): Cherry-pick into release branch
If this PR also needs to be merged into
core_r*release branches, after this PR has been merged, selectCherry-pickto open a new PR into the release branch.For MRs into `dev` branch
The proposed review process for `dev` branch is under active discussion.MRs are mergable after one approval by either
eharper@nvidia.comorzijiey@nvidia.com.Merging your PR
Any member of core-adlr and
core-nemowill be able to merge your PR.