Skip to content

[PyTorch] Move FusedAdam/FusedSGD and necessary kernels from Apex to TE#867

Merged
timmoon10 merged 7 commits into
NVIDIA:mainfrom
yaox12:fused_adam
May 30, 2024
Merged

[PyTorch] Move FusedAdam/FusedSGD and necessary kernels from Apex to TE#867
timmoon10 merged 7 commits into
NVIDIA:mainfrom
yaox12:fused_adam

Conversation

@yaox12
Copy link
Copy Markdown
Member

@yaox12 yaox12 commented May 24, 2024

Description

This PR moves FusedAdam/FusedSGD and necessary multi_tensor_* kernels from Apex to TE, as well as the tests.

Type of change

  • Documentation change (change only to the documentation, either a fix or a new content)
  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)

Changes

Please list the changes introduced in this PR:

  • Move multi_tensor_* kernels
  • Move FusedAdam
  • Move FusedSGD
  • Move tests

Checklist:

  • I have read and followed the contributing guidelines
  • The functionality is complete
  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes

@timmoon10
Copy link
Copy Markdown
Member

/te-ci pytorch

yaox12 added 4 commits May 26, 2024 18:57
Signed-off-by: Xin Yao <xiny@nvidia.com>
Signed-off-by: Xin Yao <xiny@nvidia.com>
Signed-off-by: Xin Yao <xiny@nvidia.com>
Signed-off-by: Xin Yao <xiny@nvidia.com>
Signed-off-by: Xin Yao <xiny@nvidia.com>
@timmoon10
Copy link
Copy Markdown
Member

/te-ci pytorch

@yaox12
Copy link
Copy Markdown
Member Author

yaox12 commented May 30, 2024

@timmoon10 Can I get this PR merged if it looks good to you? It's sufficient to be a drop-in replacement of Apex.FusedAdam in MCore. Other features to be migrated or implemented will come in separate PRs.

Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Copy link
Copy Markdown
Member

@timmoon10 timmoon10 left a comment

Choose a reason for hiding this comment

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

LGTM, I confirm this is mostly copy-pasted from Apex with some minor formatting changes.

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.

2 participants