Add BF16 memoment1/2 in PARTIAL_ROWWISE_ADAM TBE #2524
Closed
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
|
This pull request was exported from Phabricator. Differential Revision: D56192737 |
✅ Deploy Preview for pytorch-fbgemm-docs ready!
To edit notification comments on pull requests, go to your Netlify site configuration. |
Summary: This diff adds the backend to support multiple optimizer state types in TBE GPU training. Previously, TBE only supported FP32 optimizer states. The key changes for supporting multiple optimizer state types include: - Adding the `ArgType.PLACEHOLDER_TENSOR` tensor type in the code generation script to indicate that a tensor can have multiple types. - Updating code generation scripts to generate appropriate template types for both GPU kernels and hosts. - Modifying the TBE GPU kernel code templates to instantiate kernels for all supported types. - Updating the TBE host code template to dispatch kernels based on the types of the optimizer states. Note that this diff still supports only FP32 optimizer states. The additional optimizer state type support will be added in the subsequent diffs. Reviewed By: spcyppt Differential Revision: D56192738
Summary: This diff adds BF16 support for momentum1 and momentum2 in PARTIAL_ROWWISE_ADAM optimizer in TBE. Backend: D56192739 (this diff), D56192738, D56192736 Frontend: D56192737 Reviewed By: spcyppt Differential Revision: D56192739
Summary: This diff adds BF16 support for momentum1 and momentum2 in PARTIAL_ROWWISE_ADAM optimizer in TBE. Backend: D56192739, D56192738, D56192736 Frontend: D56192737 (this diff) Differential Revision: D56192737
|
This pull request was exported from Phabricator. Differential Revision: D56192737 |
|
This pull request has been merged in 15b4fd3. |
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.
Summary:
This diff adds BF16 support for momentum1 and momentum2 in
PARTIAL_ROWWISE_ADAM optimizer in TBE.
Backend: D56192739, D56192738, D56192736
Frontend: D56192737 (this diff)
Differential Revision: D56192737