[PyTorch] Stop storing fused weight tensor in linear modules #719
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
Stop storing fused buffers in linear modules. Signed-off-by: Tim Moon <tmoon@nvidia.com>
|
/te-ci pytorch |
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>
|
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>
|
/te-ci pytorch |
timmoon10
added
testing
Signed-off-by: Tim Moon <tmoon@nvidia.com>
|
/te-ci pytorch |
|
/te-ci pytorch |
|
The failed test is unrelated ( |
|
Anything holding up this being merged? |
|
/te-ci pytorch |
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
|
/te-ci pytorch |
|
/te-ci pytorch |
pggPL
pushed a commit
to pggPL/TransformerEngine
that referenced
this pull request
May 15, 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>
pggPL
pushed a commit
to pggPL/TransformerEngine
that referenced
this pull request
May 16, 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>
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>
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.
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.