Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Various transformer updates to improve performance #182

Merged
merged 8 commits into from
Nov 26, 2024
Merged

Various transformer updates to improve performance #182

merged 8 commits into from
Nov 26, 2024
+522 −221

Conversation

coryMosaicML
Copy link
Collaborator

@coryMosaicML coryMosaicML commented Nov 26, 2024 •
edited
Loading

This PR includes a few updates and fixes to the diffusion transformer to improve performance

Changes include:

  • Various code refactors to avoid some inefficient ops
  • Selectable backends for the attention layer
  • Option to use both DiT and MMDiT blocks
  • Grouping of transformer blocks for more efficient sharding
  • Option to add additional register tokens
  • Added torch compile
  • Remove masking in attention layers as it is an unnecessary slowdown
  • Minor fixes to the image logging callback

@coryMosaicML coryMosaicML merged commit b7e5029 into mosaicml:main Nov 26, 2024
5 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers
No reviews
Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
2 participants
@coryMosaicML @corystephenson-db