Uh oh!
There was an error while loading. Please reload this page.
[PyTorch] Stop storing fused weight tensor in linear modules - #719
Merged
Conversation
Stop storing fused buffers in linear modules. Signed-off-by: Tim Moon <tmoon@nvidia.com>
timmoon10
commented
Mar 13, 2024
MemberAuthor
/te-ci pytorch |
timmoon10
marked this pull request as draft
March 14, 2024 20:46
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>
MemberAuthor
block=te.Linear(...).to(dtype=dtype).cuda()
I see the existing implementation reallocates the fused weight tensor correctly: It is strange to me that the tests were previously passing with bit-wise exact accuracy though. I wonder if initializing the weights directly in FP16/BF16 instead of initializing in FP32 and casting resulted in tensors that were not bit-wise identical. |
Signed-off-by: Tim Moon <tmoon@nvidia.com>
timmoon10
marked this pull request as ready for review
March 15, 2024 03:42
timmoon10
commented
Mar 15, 2024
MemberAuthor
/te-ci pytorch |
Signed-off-by: Tim Moon <tmoon@nvidia.com>
timmoon10
commented
Mar 15, 2024
MemberAuthor
/te-ci pytorch |
ksivaman
commented
Mar 25, 2024
Member
/te-ci pytorch |
timmoon10
commented
Mar 29, 2024
MemberAuthor
The failed test is unrelated ( |
deepakn94
commented
Apr 10, 2024
Contributor
Anything holding up this being merged? |
ksivaman
commented
Apr 10, 2024
Member
/te-ci pytorch |
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
timmoon10
commented
Apr 16, 2024
MemberAuthor
/te-ci pytorch |
timmoon10
commented
Apr 18, 2024
MemberAuthor
/te-ci pytorch |
pggPL pushed a commit
to pggPL/TransformerEngine
that referenced
this pull request
May 23, 2024
) * Support noop concat without providing full tensor Stop storing fused buffers in linear modules. Signed-off-by: Tim Moon <tmoon@nvidia.com> * Debug noop cat func Signed-off-by: Tim Moon <tmoon@nvidia.com> * Construct TE modules in tests with correct dtypes Signed-off-by: Tim Moon <tmoon@nvidia.com> * Add tolerances to numerical tests Signed-off-by: Tim Moon <tmoon@nvidia.com> * Use plain PyTorch concat when exporting to ONNX Signed-off-by: Tim Moon <tmoon@nvidia.com> --------- Signed-off-by: Tim Moon <tmoon@nvidia.com> Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
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
Sign up for freeto join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
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.
We've encountered excessive memory usage with FSDP-like workflows because
LinearandLayerNormLinearsometimes store a tensor for fused weights. The weight params may be manipulated externally, e.g. to deallocate them or make them views into an all-gather buffer, but the modules hold on to the original buffers and prevent any memory savings. This PR updates the_noop_catutility function so that it no longer requires the full output tensor, but can figure things out by inspecting the pointers and strides. There is one functional difference:LinearandLayerNormLinearno longer make any attempt to readjust the params if they are not contiguous, but will just concatenate as needed. Besides, if FSDP is misconfigured it will just wipe out the existing buffers and create new misaligned weights in every forward pass, so there's no hope that TE can repair the situation.This is related to #570, which removed the fused weight tensors in cases without split weight params.