Skip to content

Flash attention for GPUs like in maxtext#149

Merged
entrpn merged 35 commits intoAI-Hypercomputer:mainfrom
ksikiric:kris/flux-FA
Feb 21, 2025
Merged

Flash attention for GPUs like in maxtext#149
entrpn merged 35 commits intoAI-Hypercomputer:mainfrom
ksikiric:kris/flux-FA

Conversation

@ksikiric
Copy link
Copy Markdown
Contributor

Related to #147

Adding FA support for GPUs using TransformerEngine, same as in maxtext. These changes are added on top of #147, which has been rebased on flux_lora as per #148.

Copy link
Copy Markdown
Collaborator

@entrpn entrpn left a comment

Choose a reason for hiding this comment

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

@ksikiric left some comments. Overall looks good. Last comment is to remove the training part of the code out of this PR and we can work on it on the Flux training PR.

Comment thread src/maxdiffusion/configs/base_flux_dev.yml Outdated
Comment thread src/maxdiffusion/configs/base_flux_dev.yml Outdated
Comment thread src/maxdiffusion/configs/base_flux_dev.yml Outdated
Comment thread src/maxdiffusion/configs/base_flux_dev.yml Outdated
Comment thread src/maxdiffusion/configs/base_flux_dev.yml Outdated
Comment thread src/maxdiffusion/configs/base_flux_schnell.yml Outdated
Comment thread src/maxdiffusion/configs/base_flux_schnell.yml Outdated
Comment thread src/maxdiffusion/configs/base_flux_schnell.yml Outdated
@entrpn
Copy link
Copy Markdown
Collaborator

entrpn commented Feb 14, 2025

@ksikiric tested this and looks great. Added a few comments, once resolved this should be ready to merge into main.

Please rebase with main and run ruff --fix and bash code_style.sh.

@entrpn
Copy link
Copy Markdown
Collaborator

entrpn commented Feb 19, 2025

@ksikiric let me know when you can take a look at the latest comments. This is very close to being ready! :)

@ksikiric
Copy link
Copy Markdown
Contributor Author

@entrpn I can't see any other comments than the ones I have already marked as resolved, are there any other comments that I am missing?

Comment thread src/maxdiffusion/configs/base_flux_dev.yml Outdated
Comment thread README.md Outdated
Comment thread src/maxdiffusion/configs/base_flux_schnell.yml Outdated
@entrpn
Copy link
Copy Markdown
Collaborator

entrpn commented Feb 19, 2025

@ksikiric my bad, forgot to click the button. Take a look and let me know if you see them now.

@ksikiric
Copy link
Copy Markdown
Contributor Author

ksikiric commented Feb 20, 2025

@entrpn I fixed those comments now!

@ksikiric ksikiric marked this pull request as ready for review February 20, 2025 09:09
Comment thread src/maxdiffusion/configs/base_flux_dev.yml Outdated
@entrpn entrpn merged commit 1c9d4c1 into AI-Hypercomputer:main Feb 21, 2025
@entrpn
Copy link
Copy Markdown
Collaborator

entrpn commented Feb 21, 2025

@entrpn I fixed those comments now!

Thank you.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants