Cherry picking keep_fp8_weight_transpose_cache flag refactor and fsdp2 fp8 autocast all gather commits #389
Conversation
This file contains hidden or 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
* Initial commit * Removed rocm_utils * Added comment and bug fixes * Grouped IS_HIP_EXTENSION with the property assignment * Reverted transpose.cpp, removed keep_fp8_transpose_cache flag from grouped_linear, removed manual clearing of tensors in modules * Aligning grouped_linear module with upstream * Reverted tests to use _test_granular_accuracy_with_fp8 multiple times as needed * Added comments back * Moved comment to the test --------- Co-authored-by: sudhu2k <[email protected]>
sudhu2k
requested review from
ipanfilo,
wangye805 and
wenchenvincent
as code owners
* Initial commit * Removed Print statements, added keep_fp8_transpose cache integration with fsdp2 * Added use_fsdp flag to Linear module, added profile code, added test code, added all reduce for amax * Fixed unit test * Removing all reduce code for amax since by default TE does all reduce when torch.distributed is initialized. * reverting case where out is already present * Added unit test with regualr sgpu training * Modified unit test to compare FSDP2 with DDP * bug fixes * Code cleaning up * Initial commit to add MXFP8 * Added fp8 current scaling. * Added MXFP8, Modified unit test to run based on recipes * Extended use_fsdp to layernorm linear and layernorm mlp * Moved amax reduce from forward to backward for fsdp2 * Added automatic detection of use fsdp from base module * Use SKIP_FP8_REDUCTION_FOR_FSDP2 in backward for check if need to do forward reduce * Added memory profile code, added a check before setting SKIP_FP8_REDUCTION_FOR_FSDP2 * Fix for fused optimizer, changed _elem to _data, code clean up * Fixed layernorm mlp * Code cleanup and added test to pytorch.sh * Removed whitespaces * Fixed comments and license * Added guards * Added reduce for forward in cuda graph backward, added code to remove test artifacts, reverted upstream test file --------- Co-authored-by: sudhu2k <[email protected]>
sudhu2k
changed the title
Contributor
Author
|
@ipanfilo, @wangye805 The PR is ready for review |
ipanfilo
requested changes
Dec 3, 2025
Contributor
Author
There was a problem hiding this comment.
But nothing significant was added from our side, this PR actually removes the code that was added for keep_fp8_weight_transpose_cache which means technically we are reverting to upstream code for grouped_linear.py
Collaborator
There was a problem hiding this comment.
Does it match upstream now?
ipanfilo
approved these changes
Dec 3, 2025
Collaborator
ipanfilo
left a comment
ipanfilo
left a comment
There was a problem hiding this comment.
Let's wait for CI to be run and passed
…spose-refactor-cherrypick-rv2.2
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.
Description
This PR is used for cherry picking #349 and #328 into release_v2.4 branch
Type of change
Checklist: