[JAX] Fix layernorm distributed test sharding and collective assertions #2041
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
Signed-off-by: Jeremy Berchtold <[email protected]>
Signed-off-by: Jeremy Berchtold <[email protected]>
jberchtold-nvidia
force-pushed
the
dev/jberchtold/fix-layernorm-distributed-tests
branch
from
cc8be88 to
3fef659
Compare
jberchtold-nvidia
changed the title
|
/te-ci L1 jax |
Signed-off-by: Jeremy Berchtold <[email protected]>
jberchtold-nvidia
force-pushed
the
dev/jberchtold/fix-layernorm-distributed-tests
branch
from
a5b8244 to
025db08
Compare
Signed-off-by: Jeremy Berchtold <[email protected]>
jberchtold-nvidia
force-pushed
the
dev/jberchtold/fix-layernorm-distributed-tests
branch
from
2006c1f to
e3d2a48
Compare
|
/te-ci L1 jax |
Signed-off-by: Jeremy Berchtold <[email protected]>
jberchtold-nvidia
force-pushed
the
dev/jberchtold/fix-layernorm-distributed-tests
branch
from
17ae015 to
629dba5
Compare
|
/te-ci L1 jax |
jberchtold-nvidia
force-pushed
the
dev/jberchtold/fix-layernorm-distributed-tests
branch
from
a3a9513 to
629dba5
Compare
|
/te-ci L1 jax |
1 similar comment
|
/te-ci L1 jax |
phu0ngng
reviewed
Aug 11, 2025
| @@ -828,18 +828,23 @@ def _jax_layernorm(x, gamma, beta, zero_centered_gamma, epsilon, quantizer=None) | |||
| """ | |||
| JAX native layernorm implementation | |||
| """ | |||
| x_ = jnp.asarray(x, jnp.float32) | |||
There was a problem hiding this comment.
I'm not sure about this change, as without this casting, the ctype of normalization will be bf16.
| f"{name}={'true' if enabled else 'false'}" for name, enabled in primitives.items() | ||
| ) | ||
|
|
||
| os.environ["NVTE_JAX_CUSTOM_CALLS"] = str(updated_env_var) |
There was a problem hiding this comment.
Hi,
I think we could use manage_primitives here and don't need to change the env var.
If you prefer to pass key-value pairs, we could update the manage_primitives to do so too.
There was a problem hiding this comment.
Thanks, just read the full docstring on manage_primitives and you're right it will work. I'll update to that instead for this context
I had missed the disable_all_first argument and thought manage_primitives was always disabling any unspecified primitives
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
Disables TE norm primitive. Updates distributed layernorm collection communication byte count logic and weight sharding. Also extend tests to run both TE norm and JAX norm
Type of change
Changes
tests/jax/test_distributed_layernorm.mlpprimitive_contextfor modifying which TE primitives are enabled within a blockuse_jax_gemmto use this utility instead. Slight improvement on existing behavior as previouslyuse_jax_gemmwould ignoreNVTE_JAX_CUSTOM_CALLSwhen inside the block, but now the behavior is(all updates specified in NVTE_JAX_CUSTOM_CALLS) + (override GemmPrimitive=true/false)Checklist: