[JAX] Fix partitioning issues in LayerNorm and LayerNormMLP layers #1743
+28
−7
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
jberchtold-nvidia
changed the title
Signed-off-by: Jeremy Berchtold <[email protected]>
…hape checks before calling TE API Signed-off-by: Jeremy Berchtold <[email protected]>
jberchtold-nvidia
force-pushed
the
dev/jberchtold/fix-issues-for-maxtext-integration
branch
from
5350e48 to
1260f07
Compare
|
/te-ci L0 |
phu0ngng
changed the title
phu0ngng
reviewed
May 5, 2025
Comment on lines
278
to
282
| NVTE_CHECK(act_len == 1 || act_len == 2, | ||
| "The value of the activation dimension (axis=-2) must be 1 for non-gated or 2 for " | ||
| "gated activation, got ", | ||
| act_len); | ||
| checkDActShapes(input_buf, act_input_buf, output_buf); |
There was a problem hiding this comment.
I think we should be able to check all of these requirements in the Primitive, in the abstract() or/and lowering().
There was a problem hiding this comment.
Good point, I guess checking in the abstract will enforce the full shape and the sharded shape as we call the abstract again on the inner sharded primitive
There was a problem hiding this comment.
Moved this check to the abstract function in python and confirmed it still catches the issue when the partitioning fix isn't applied
Signed-off-by: Jeremy Berchtold <[email protected]>
|
/te-ci L0 |
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
Fixes two issues occurring in MaxText integration with TE. 1) NaNs caused on multi-gpu usage of the LayerNormMLP layer 2) Partitioning shape mismatch errors caused by using TE's LayerNorm layer with TP>1
Type of change
Changes
dzandxare the same shape.dzalways use the same partitioning asx. In combination with the checks above, this prevents NaNs from running dact ondzandxpartitioned differently leading to out of bounds reads.xnot partition along the TP axis. This prevents partitioning shape errors we were encountering wherexwas partitioned along TP but the norm output wasn't.Checklist: