Repository: NVIDIA/TransformerEngine
Pull request #1772: [JAX] Support logical partitioning axes in TE Flax modules
Author: login: jberchtold-nvidia
Labels: name: 2.4.0
Created: 2025-05-12T16:17:01Z | Merged: 2025-05-16T17:34:15Z

# Description

TE/JAX Flax modules are using weight initialization that only supports physical mesh partitioning axes, not logical axes used in some frameworks, such as MaxText. To support partitioning in these frameworks weight initialization was updated to a newer way that supports logical axes.

## Type of change

- [ ] Documentation change (change only to the documentation, either a fix or a new content)
- [x] Bug fix (non-breaking change which fixes an issue)
- [ ] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected)
- [ ] Infra/Build change
- [ ] Code refactoring

## Changes

Please list the changes introduced in this PR:

- Update Flax modules in `module.py` and `transformer.py` to replace usages of `nn_partitioning.param_with_axes` with `self.param(..., nn.with_logical_partitioning(...), ...)`
- Fix bug where ffn1 intermediate result in the LayerNormMLP layer incorrectly added a sharding constraint to replicate the tensor and following gemm, leading to OOM on some configs that normal MaxText did not have OOM on.

# Checklist:

- [x] I have read and followed the [contributing guidelines](https://github.com/NVIDIA/TransformerEngine/blob/main/CONTRIBUTING.rst)
- [x] The functionality is complete
- [x] I have commented my code, particularly in hard-to-understand areas
- [x] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my feature works
- [x] New and existing unit tests pass locally with my changes



Diff:
diff --git a/examples/jax/encoder/test_multiprocessing_encoder.py b/examples/jax/encoder/test_multiprocessing_encoder.py
index 2cf56cdf37..ec324d7b37 100644
--- a/examples/jax/encoder/test_multiprocessing_encoder.py
+++ b/examples/jax/encoder/test_multiprocessing_encoder.py
@@ -609,7 +609,7 @@ def test_te_bf16(self):
     def test_te_delayed_scaling_fp8(self):
         """Test Transformer Engine with DelayedScaling FP8"""
         result = self.exec(True, "DelayedScaling")
-        assert result[0] < 0.505 and result[1] > 0.754
+        assert result[0] < 0.505 and result[1] > 0.753
 
     @unittest.skipIf(
         not is_fp8_supported(), "Device compute capability 9.0+ is required for CurrentScaling FP8"
diff --git a/tests/jax/utils.py b/tests/jax/utils.py
index 24321a5020..39eb368375 100644
--- a/tests/jax/utils.py
+++ b/tests/jax/utils.py
@@ -13,7 +13,6 @@
 import jax.numpy as jnp
 import numpy as np
 from flax import linen as nn
-from flax.linen import partitioning as nn_partitioning
 from flax.linen.attention import combine_masks
 from jax import lax, vmap
 from jax import nn as jax_nn
@@ -316,16 +315,22 @@ def __call__(self, inputs: Array) -> Array:
 
         kernel_shape = tuple(inputs.shape[ax] for ax in axis) + features
         kernel_param_shape = (np.prod([inputs.shape[ax] for ax in axis]), np.prod(features))
-        kernel = nn_partitioning.param_with_axes(
-            "kernel", self.kernel_init, kernel_param_shape, self.dtype, axes=self.kernel_axes
+        kernel = self.param(
+            "kernel",
+            nn.with_logical_partitioning(self.kernel_init, self.kernel_axes),
+            kernel_param_shape,
+            self.dtype,
         )
 
         kernel = jnp.asarray(kernel, input_dtype)
         kernel = jnp.reshape(kernel, kernel_shape)
 
         if self.use_bias:
-            bias = nn_partitioning.param_with_axes(
-                "bias", self.bias_init, self.features, self.dtype, axes=self.bias_axes
+            bias = self.param(
+                "bias",
+                nn.with_logical_partitioning(self.bias_init, self.bias_axes),
+                self.features,
+                self.dtype,
             )
             bias = bias.astype(input_dtype)
         else:
@@ -422,9 +427,9 @@ def __call__(self, inputs, deterministic: bool = False):
         )  # Broadcast along length.
 
         if self.transpose_batch_sequence:
-            x = nn_partitioning.with_sharding_constraint(x, ("length", "batch", "mlp"))
+            x = nn.with_logical_constraint(x, ("length", "batch", "mlp"))
         else:
-            x = nn_partitioning.with_sharding_constraint(x, ("batch", "length", "mlp"))
+            x = nn.with_logical_constraint(x, ("batch", "length", "mlp"))
         output = DenseGeneral(
             inputs.shape[-1],
             dtype=self.dtype,
@@ -688,21 +693,13 @@ def qkv_init(key, shape, dtype):
         value = value.reshape((*value.shape[:2], self.num_gqa_groups, self.head_dim))
 
         if self.transpose_batch_sequence:
-            query = nn_partitioning.with_sharding_constraint(
-                query, ("length", "batch", "heads", "kv")
-            )
-            key = nn_partitioning.with_sharding_constraint(key, ("length", "batch", "heads", "kv"))
-            value = nn_partitioning.with_sharding_constraint(
-                value, ("length", "batch", "heads", "kv")
-            )
+            query = nn.with_logical_constraint(query, ("length", "batch", "heads", "kv"))
+            key = nn.with_logical_constraint(key, ("length", "batch", "heads", "kv"))
+            value = nn.with_logical_constraint(value, ("length", "batch", "heads", "kv"))
         else:
-            query = nn_partitioning.with_sharding_constraint(
-                query, ("batch", "length", "heads", "kv")
-            )
-            key = nn_partitioning.with_sharding_constraint(key, ("batch", "length", "heads", "kv"))
-            value = nn_partitioning.with_sharding_constraint(
-                value, ("batch", "length", "heads", "kv")
-            )
+            query = nn.with_logical_constraint(query, ("batch", "length", "heads", "kv"))
+            key = nn.with_logical_constraint(key, ("batch", "length", "heads", "kv"))
+            value = nn.with_logical_constraint(value, ("batch", "length", "heads", "kv"))
 
         if decode:
             # Detect if we're initializing by absence of existing cache data.
@@ -809,9 +806,9 @@ def qkv_init(key, shape, dtype):
         x = x.reshape((x.shape[0], x.shape[1], x.shape[2] * x.shape[3]))
 
         if self.transpose_batch_sequence:
-            x = nn_partitioning.with_sharding_constraint(x, ("length", "batch", "joined_kv"))
+            x = nn.with_logical_constraint(x, ("length", "batch", "joined_kv"))
         else:
-            x = nn_partitioning.with_sharding_constraint(x, ("batch", "length", "joined_kv"))
+            x = nn.with_logical_constraint(x, ("batch", "length", "joined_kv"))
 
         # Back to the original inputs dimensions.
 
@@ -857,8 +854,11 @@ def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
         input_dtype = x.dtype
         features = x.shape[-1]
 
-        scale = nn_partitioning.param_with_axes(
-            "scale", self.scale_init, (features,), self.dtype, axes=("embed",)
+        scale = self.param(
+            "scale",
+            nn.with_logical_partitioning(self.scale_init, ("embed",)),
+            (features,),
+            self.dtype,
         )
         x_ = x.astype(jnp.float32)
         if self.layernorm_type == "layernorm":
@@ -866,8 +866,11 @@ def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
             var = jnp.mean(jnp.square(x_ - mean), axis=-1, keepdims=True)
             y = (x_ - mean) * lax.rsqrt(var + self.epsilon)
 
-            bias = nn_partitioning.param_with_axes(
-                "ln_bias", self.bias_init, (features,), self.dtype, axes=("embed",)
+            bias = self.param(
+                "ln_bias",
+                nn.with_logical_partitioning(self.bias_init, ("embed",)),
+                (features,),
+                self.dtype,
             )
             bias = jnp.asarray(bias, input_dtype)
 
@@ -976,12 +979,11 @@ def __call__(self, qlen, klen, bidirectional=True):
             num_buckets=self.num_buckets,
             max_distance=self.max_distance,
         )
-        relative_attention_bias = nn_partitioning.param_with_axes(
+        relative_attention_bias = self.param(
             "rel_embedding",
-            self.embedding_init,
+            nn.with_logical_partitioning(self.embedding_init, ("heads", "relpos_buckets")),
             (self.num_heads, self.num_buckets),
             jnp.float32,
-            axes=("heads", "relpos_buckets"),
         )
 
         relative_attention_bias = jnp.asarray(relative_attention_bias, self.dtype)
@@ -1559,14 +1561,16 @@ def sync_params_values(dst, src, transformations, sep="/"):
     """
     src_values = {}
     for key, value in jax.tree_util.tree_leaves_with_path(src):
-        normalized_key = sep.join(x.key for x in key)
+        # Only select DictKey(key="...") entries, skip GetAttr(name="...") entries at the end of the tree path
+        normalized_key = sep.join(x.key for x in key if hasattr(x, "key"))
         src_values[normalized_key] = value
 
     flatten_dst, dst_tree_def = jax.tree_util.tree_flatten_with_path(dst)
     synced_dst_values = []
 
     for key, value in flatten_dst:
-        normalized_key = sep.join(x.key for x in key)
+        # Only select DictKey(key="...") entries, skip GetAttr(name="...") entries at the end of the tree path
+        normalized_key = sep.join(x.key for x in key if hasattr(x, "key"))
         if normalized_key in transformations:
             corresponding_src_key = transformations[normalized_key]
         else:
diff --git a/transformer_engine/jax/dense.py b/transformer_engine/jax/dense.py
index f8d053afb1..55d60e4189 100644
--- a/transformer_engine/jax/dense.py
+++ b/transformer_engine/jax/dense.py
@@ -49,6 +49,7 @@ def dense(
     """
     # Remove when tex.quantize() can handle quantizer=None
     if quantizer_set == noop_quantizer_set:
+        x = with_sharding_constraint_by_logical_axes(x, input_axes)
         output = tex.gemm(x, kernel, contracting_dims)
         if bias is not None:
             bias_new_shape = (1,) * (output.ndim - bias.ndim) + bias.shape
diff --git a/transformer_engine/jax/flax/module.py b/transformer_engine/jax/flax/module.py
index ef60052768..bd311472f0 100644
--- a/transformer_engine/jax/flax/module.py
+++ b/transformer_engine/jax/flax/module.py
@@ -11,7 +11,6 @@
 import numpy as np
 import jax.numpy as jnp
 from flax import linen as nn
-from flax.linen import partitioning as nn_partitioning
 from jax import lax
 from jax import random as jax_random
 from jax.ad_checkpoint import checkpoint_name
@@ -65,6 +64,7 @@ def _obtain_default_layernorm_scale_init_if_need(original_init, zero_centered_ga
 
 
 def _create_layernorm_parameters(
+    module,
     norm_type,
     shape,
     scale_init,
@@ -74,13 +74,21 @@ def _create_layernorm_parameters(
     input_dtype,
     dtype,
 ):
-    scale = nn_partitioning.param_with_axes("scale", scale_init, shape, dtype, axes=scale_axes)
-    scale = scale.astype(input_dtype)
+    scale = module.param(
+        "scale",
+        nn.with_logical_partitioning(scale_init, scale_axes),
+        shape,
+        dtype,
+    ).astype(input_dtype)
 
     norm_type = canonicalize_norm_type(norm_type)
     if norm_type == "layernorm":
-        bias = nn_partitioning.param_with_axes("ln_bias", bias_init, shape, dtype, axes=bias_axes)
-        bias = jnp.asarray(bias, input_dtype)
+        bias = module.param(
+            "ln_bias",
+            nn.with_logical_partitioning(bias_init, bias_axes),
+            shape,
+            dtype,
+        ).astype(input_dtype)
     else:
         assert norm_type == "rmsnorm"
         bias = None
@@ -308,6 +316,7 @@ def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
 
         features = x.shape[-1]
         scale, ln_bias = _create_layernorm_parameters(
+            self,
             self.layernorm_type,
             (features,),
             self.scale_init,
@@ -467,16 +476,22 @@ def __call__(self, inputs: Array) -> Array:
                 "Expected len(kernel_shape) to match len(kernel_axes),"
                 f"got kernel_shape {kernel_shape} and kernel_axes {self.kernel_axes}"
             )
-        kernel = nn_partitioning.param_with_axes(
-            "kernel", self.kernel_init, kernel_shape, self.dtype, axes=self.kernel_axes
+        kernel = self.param(
+            "kernel",
+            nn.with_logical_partitioning(self.kernel_init, self.kernel_axes),
+            kernel_shape,
+            self.dtype,
         )
 
         if not QuantizeConfig.is_fp8_enabled():
             kernel = kernel.astype(input_dtype)
 
         if self.use_bias:
-            bias = nn_partitioning.param_with_axes(
-                "bias", self.bias_init, features, self.dtype, axes=self.bias_axes
+            bias = self.param(
+                "bias",
+                nn.with_logical_partitioning(self.bias_init, self.bias_axes),
+                features,
+                self.dtype,
             ).astype(input_dtype)
         else:
             bias = None
@@ -499,25 +514,21 @@ def __call__(self, inputs: Array) -> Array:
                 self.low_rank_adaptation_dim,
             )
             lora_a_kernel_axes = (None,) * len(lora_a_kernel_shape)
-            lora_a_kernel = nn_partitioning.param_with_axes(
+            lora_a_kernel = self.param(
                 "lora_a_kernel",
-                self.kernel_init,
+                nn.with_logical_partitioning(self.kernel_init, lora_a_kernel_axes),
                 lora_a_kernel_shape,
                 self.dtype,
-                axes=lora_a_kernel_axes,
-            )
-            lora_a_kernel = lora_a_kernel.astype(input_dtype)
+            ).astype(input_dtype)
 
             lora_b_kernel_shape = (*features[:-1], self.low_rank_adaptation_dim, features[-1])
             lora_b_kernel_axes = (None,) * len(lora_b_kernel_shape)
-            lora_b_kernel = nn_partitioning.param_with_axes(
+            lora_b_kernel = self.param(
                 "lora_b_kernel",
-                nn.initializers.zeros,
+                nn.with_logical_partitioning(nn.initializers.zeros, lora_b_kernel_axes),
                 lora_b_kernel_shape,
                 self.dtype,
-                axes=lora_b_kernel_axes,
-            )
-            lora_b_kernel = lora_b_kernel.astype(input_dtype)
+            ).astype(input_dtype)
 
             y += _apply_low_rank_adaptation(
                 inputs, axis, features, lora_a_kernel, lora_b_kernel, self.low_rank_adaptation_alpha
@@ -695,6 +706,7 @@ def __call__(self, inputs: Array) -> Array:
             inputs = with_sharding_constraint_by_logical_axes(inputs, self.layernorm_input_axes)
             features = inputs.shape[-1]
             scale, ln_bias = _create_layernorm_parameters(
+                self,
                 self.layernorm_type,
                 (features,),
                 self.scale_init,
@@ -730,8 +742,11 @@ def __call__(self, inputs: Array) -> Array:
         axis = _normalize_axes(axis, y.ndim)
 
         kernel_shape = (np.prod([inputs.shape[ax] for ax in axis]),) + features
-        kernel = nn_partitioning.param_with_axes(
-            "kernel", self.kernel_init, kernel_shape, self.dtype, axes=self.kernel_axes
+        kernel = self.param(
+            "kernel",
+            nn.with_logical_partitioning(self.kernel_init, self.kernel_axes),
+            kernel_shape,
+            self.dtype,
         )
         if not QuantizeConfig.is_fp8_enabled():
             kernel = kernel.astype(input_dtype)
@@ -770,25 +785,21 @@ def __call__(self, inputs: Array) -> Array:
                 self.low_rank_adaptation_dim,
             )
             lora_a_kernel_axes = (None,) * len(lora_a_kernel_shape)
-            lora_a_kernel = nn_partitioning.param_with_axes(
+            lora_a_kernel = self.param(
                 "lora_a_kernel",
-                self.kernel_init,
+                nn.with_logical_partitioning(self.kernel_init, lora_a_kernel_axes),
                 lora_a_kernel_shape,
                 self.dtype,
-                axes=lora_a_kernel_axes,
-            )
-            lora_a_kernel = lora_a_kernel.astype(input_dtype)
+            ).astype(input_dtype)
 
             lora_b_kernel_shape = (*features[:-1], self.low_rank_adaptation_dim, features[-1])
             lora_b_kernel_axes = (None,) * len(lora_b_kernel_shape)
-            lora_b_kernel = nn_partitioning.param_with_axes(
+            lora_b_kernel = self.param(
                 "lora_b_kernel",
-                nn.initializers.zeros,
+                nn.with_logical_partitioning(nn.initializers.zeros, lora_b_kernel_axes),
                 lora_b_kernel_shape,
                 self.dtype,
-                axes=lora_b_kernel_axes,
-            )
-            lora_b_kernel = lora_b_kernel.astype(input_dtype)
+            ).astype(input_dtype)
 
             z += _apply_low_rank_adaptation(
                 y, axis, features, lora_a_kernel, lora_b_kernel, self.low_rank_adaptation_alpha
@@ -796,8 +807,11 @@ def __call__(self, inputs: Array) -> Array:
 
         bias = None
         if self.use_bias:
-            bias = nn_partitioning.param_with_axes(
-                "bias", self.bias_init, features, self.dtype, axes=self.bias_axes
+            bias = self.param(
+                "bias",
+                nn.with_logical_partitioning(self.bias_init, self.bias_axes),
+                features,
+                self.dtype,
             ).astype(input_dtype)
 
         if bias is not None:
@@ -1028,6 +1042,7 @@ def __call__(self, inputs: Array, deterministic: bool = False) -> Array:
             features = inputs.shape[-1]
 
             scale, ln_bias = _create_layernorm_parameters(
+                self,
                 self.layernorm_type,
                 (features,),
                 self.scale_init,
@@ -1067,14 +1082,13 @@ def kernel_1_init(key, num_kernels, stack_axis, *init_args):
         axis = _canonicalize_tuple(self.axis)
         axis = _normalize_axes(axis, y.ndim)
         kernel_1_each_shape = (np.prod([inputs.shape[ax] for ax in axis]), self.intermediate_dim)
-        kernel_1 = nn_partitioning.param_with_axes(
+        kernel_1 = self.param(
             "wi_kernel",
-            kernel_1_init,
+            nn.with_logical_partitioning(kernel_1_init, self.kernel_axes_1),
             num_activations,
             -2,
             kernel_1_each_shape,
             self.dtype,
-            axes=self.kernel_axes_1,
         )
 
         if not QuantizeConfig.is_fp8_enabled():
@@ -1083,12 +1097,11 @@ def kernel_1_init(key, num_kernels, stack_axis, *init_args):
         hidden_size = inputs.shape[-1]
         hidden_size_tuple = _canonicalize_tuple(hidden_size)
         kernel_2_shape = (self.intermediate_dim,) + hidden_size_tuple
-        kernel_2 = nn_partitioning.param_with_axes(
+        kernel_2 = self.param(
             "wo_kernel",
-            self.kernel_init,
+            nn.with_logical_partitioning(self.kernel_init, self.kernel_axes_2),
             kernel_2_shape,
             self.dtype,
-            axes=self.kernel_axes_2,
         )
         if not QuantizeConfig.is_fp8_enabled():
             kernel_2 = kernel_2.astype(input_dtype)
@@ -1097,21 +1110,19 @@ def kernel_1_init(key, num_kernels, stack_axis, *init_args):
 
         if self.use_bias:
             bias_1_shape = (num_activations, self.intermediate_dim)
-            bias_1 = nn_partitioning.param_with_axes(
+            bias_1 = self.param(
                 "wi_bias",
-                self.bias_init,
+                nn.with_logical_partitioning(self.bias_init, self.bias_axes_1),
                 bias_1_shape,
                 self.dtype,
-                axes=self.bias_axes_1,
             ).astype(input_dtype)
 
             bias_2_shape = (hidden_size,)
-            bias_2 = nn_partitioning.param_with_axes(
+            bias_2 = self.param(
                 "wo_bias",
-                self.bias_init,
+                nn.with_logical_partitioning(self.bias_init, self.bias_axes_2),
                 bias_2_shape,
                 self.dtype,
-                axes=self.bias_axes_2,
             ).astype(input_dtype)
         else:
             bias_1 = None
@@ -1168,11 +1179,15 @@ def kernel_1_init(key, num_kernels, stack_axis, *init_args):
                     kernel_axes=self.kernel_axes_1,
                     quantizer_set=ffn1_quantizer_set,
                 )
-            dot_1_output_axes = (
-                *get_non_contracting_logical_axes(y.ndim, self.dot_1_input_axes, axis),
-                *get_non_contracting_logical_axes(kernel_1.ndim, self.kernel_axes_1, contract_ind),
-            )
-            x = with_sharding_constraint_by_logical_axes(x, dot_1_output_axes)
+
+            if self.dot_1_input_axes is not None and self.kernel_axes_1 is not None:
+                dot_1_output_axes = (
+                    *get_non_contracting_logical_axes(y.ndim, self.dot_1_input_axes, axis),
+                    *get_non_contracting_logical_axes(
+                        kernel_1.ndim, self.kernel_axes_1, contract_ind
+                    ),
+                )
+                x = with_sharding_constraint_by_logical_axes(x, dot_1_output_axes)
 
             if self.enable_low_rank_adaptation:
                 wi_lora_a_kernel_each_shape = (
@@ -1180,16 +1195,14 @@ def kernel_1_init(key, num_kernels, stack_axis, *init_args):
                     self.low_rank_adaptation_dim,
                 )
                 wi_lora_a_kernel_axes = (None,) * len(wi_lora_a_kernel_each_shape + 1)
-                wi_lora_a_kernel = nn_partitioning.param_with_axes(
+                wi_lora_a_kernel = self.param(
                     "wi_lora_a_kernel",
-                    kernel_1_init,
+                    nn.with_logical_partitioning(kernel_1_init, wi_lora_a_kernel_axes),
                     num_activations,
                     -2,
                     wi_lora_a_kernel_each_shape,
                     self.dtype,
-                    axes=wi_lora_a_kernel_axes,
-                )
-                wi_lora_a_kernel = wi_lora_a_kernel.astype(input_dtype)
+                ).astype(input_dtype)
 
                 wi_lora_b_kernel_shape = (
                     num_activations,
@@ -1197,14 +1210,12 @@ def kernel_1_init(key, num_kernels, stack_axis, *init_args):
                     self.intermediate_dim,
                 )
                 wi_lora_b_kernel_axes = (None,) * len(wi_lora_b_kernel_shape)
-                wi_lora_b_kernel = nn_partitioning.param_with_axes(
+                wi_lora_b_kernel = self.param(
                     "wi_lora_b_kernel",
-                    nn.initializers.zeros,
+                    nn.with_logical_partitioning(nn.initializers.zeros, wi_lora_b_kernel_axes),
                     wi_lora_b_kernel_shape,
                     self.dtype,
-                    axes=wi_lora_b_kernel_axes,
-                )
-                wi_lora_b_kernel = wi_lora_b_kernel.astype(input_dtype)
+                ).astype(input_dtype)
 
                 x += _apply_low_rank_adaptation(
                     y,
@@ -1253,25 +1264,21 @@ def kernel_1_init(key, num_kernels, stack_axis, *init_args):
             if self.enable_low_rank_adaptation:
                 wo_lora_a_kernel_shape = (self.intermediate_dim, self.low_rank_adaptation_dim)
                 wo_lora_a_kernel_axes = (None,) * len(wo_lora_a_kernel_shape)
-                wo_lora_a_kernel = nn_partitioning.param_with_axes(
+                wo_lora_a_kernel = self.param(
                     "wo_lora_a_kernel",
-                    self.kernel_init,
+                    nn.with_logical_partitioning(self.kernel_init, wo_lora_a_kernel_axes),
                     wo_lora_a_kernel_shape,
                     self.dtype,
-                    axes=wo_lora_a_kernel_axes,
-                )
-                wo_lora_a_kernel = wo_lora_a_kernel.astype(input_dtype)
+                ).astype(input_dtype)
 
                 wo_lora_b_kernel_shape = (self.low_rank_adaptation_dim, hidden_size)
                 wo_lora_b_kernel_axes = (None,) * len(wo_lora_b_kernel_shape)
-                wo_lora_b_kernel = nn_partitioning.param_with_axes(
+                wo_lora_b_kernel = self.param(
                     "wo_lora_b_kernel",
-                    nn.initializers.zeros,
+                    nn.with_logical_partitioning(nn.initializers.zeros, wo_lora_b_kernel_axes),
                     wo_lora_b_kernel_shape,
                     self.dtype,
-                    axes=wo_lora_b_kernel_axes,
-                )
-                wo_lora_b_kernel = wo_lora_b_kernel.astype(input_dtype)
+                ).astype(input_dtype)
 
                 out += _apply_low_rank_adaptation(
                     z,
diff --git a/transformer_engine/jax/flax/transformer.py b/transformer_engine/jax/flax/transformer.py
index 10a0a06824..1ac22a6d2f 100644
--- a/transformer_engine/jax/flax/transformer.py
+++ b/transformer_engine/jax/flax/transformer.py
@@ -15,7 +15,6 @@
 import jax.numpy as jnp
 import numpy as np
 from flax import linen as nn
-from flax.linen import partitioning as nn_partitioning
 from flax.linen.attention import combine_masks
 from jax import nn as jax_nn
 from jax import random as jax_random
@@ -1503,12 +1502,11 @@ def __call__(self, q_seqlen, k_seqlen, bidirectional=True):
         rp_bucket += np.where(rpb_is_small, negative_rp, rpb_val_if_large)
 
         # Compute relative attention bias
-        relative_attention_bias = nn_partitioning.param_with_axes(
+        relative_attention_bias = self.param(
             "rel_embedding",
-            self.embedding_init,
+            nn.with_logical_partitioning(self.embedding_init, self.embedding_axes),
             (self.num_attention_heads, self.num_buckets),
             self.dtype,
-            axes=self.embedding_axes,
         )
 
         relative_attention_bias = jnp.asarray(relative_attention_bias, self.dtype)
diff --git a/transformer_engine/jax/layernorm_mlp.py b/transformer_engine/jax/layernorm_mlp.py
index e7e3fd2fb9..e04b930233 100644
--- a/transformer_engine/jax/layernorm_mlp.py
+++ b/transformer_engine/jax/layernorm_mlp.py
@@ -275,11 +275,12 @@ def _layernorm_mlp_fwd_rule(
         (x_contracting_dims, k_contracting_dims),
     )
 
-    dot_1_output_axes = (
-        *get_non_contracting_logical_axes(x.ndim, dot_1_input_axes, x_contracting_dims),
-        *get_non_contracting_logical_axes(kernel_1.ndim, kernel_1_axes, k_contracting_dims),
-    )
-    dot_1_output = with_sharding_constraint_by_logical_axes(dot_1_output, dot_1_output_axes)
+    if dot_1_input_axes is not None and kernel_1_axes is not None:
+        dot_1_output_axes = (
+            *get_non_contracting_logical_axes(x.ndim, dot_1_input_axes, x_contracting_dims),
+            *get_non_contracting_logical_axes(kernel_1.ndim, kernel_1_axes, k_contracting_dims),
+        )
+        dot_1_output = with_sharding_constraint_by_logical_axes(dot_1_output, dot_1_output_axes)
 
     if use_bias_1:
         bias_1_shape = bias_1.shape
@@ -303,12 +304,6 @@ def _layernorm_mlp_fwd_rule(
         (x_contracting_dims, k_contracting_dims),
     )
 
-    dot_2_output_axes = (
-        *get_non_contracting_logical_axes(x.ndim, dot_2_input_axes, x_contracting_dims),
-        *get_non_contracting_logical_axes(kernel_2.ndim, None, k_contracting_dims),
-    )
-    dot_2_output = with_sharding_constraint_by_logical_axes(dot_2_output, dot_2_output_axes)
-
     if use_bias_2:
         bias_2_shape = bias_2.shape
         bias_2_new_shape = (1,) * (dot_2_output.ndim - bias_2.ndim) + bias_2_shape
diff --git a/transformer_engine/jax/sharding.py b/transformer_engine/jax/sharding.py
index df3f38cbd1..f89fe60f95 100644
--- a/transformer_engine/jax/sharding.py
+++ b/transformer_engine/jax/sharding.py
@@ -13,7 +13,7 @@
 from contextlib import contextmanager
 from dataclasses import dataclass
 from enum import Enum
-from typing import Callable
+from typing import Callable, Optional
 from jax.interpreters import pxla
 import jax
 import jax.numpy as jnp
@@ -112,9 +112,22 @@ def with_sharding_constraint(x: jnp.array, pspec: PartitionSpec):
     return jax.lax.with_sharding_constraint(x, pspec)
 
 
-def with_sharding_constraint_by_logical_axes(x: jnp.array, logical_axis_names: tuple | list):
+def with_sharding_constraint_by_logical_axes(
+    x: jnp.array, logical_axis_names: Optional[tuple | list]
+):
     """
     A wrapper function to jax.lax.with_sharding_constraint to accept logical axes.
+
+    If logical_axis_names = None, this means no sharding constraint is applied.
+
+    If logical_axis_names = (None, None, ...), this means a sharding constraint is applied and the tensor is replicated across all devices.
+
+    Args:
+        x: Input tensor to apply sharding constraint
+        logical_axis_names: Logical axis names to apply sharding constraint
+    Returns:
+        Tensor with sharding constraint applied, or the original tensor if no logical axes are provided.
+
     """
     if not logical_axis_names:
         return x
@@ -321,7 +334,9 @@ class ShardingType(Enum):
     DP_TP_ROW = (MajorShardingType.DPTP, "dp_tp_row")
 
 
-def get_non_contracting_logical_axes(ndim, logical_axes, contracting_dims):
+def get_non_contracting_logical_axes(
+    ndim, logical_axes: tuple[Optional[str]], contracting_dims
+) -> tuple[Optional[str]]:
     """Get logical axes for non-contracting dimensions.
 
     Args:
@@ -332,11 +347,8 @@ def get_non_contracting_logical_axes(ndim, logical_axes, contracting_dims):
     Returns:
         Tuple of logical axes for non-contracting dimensions.
     """
-    if not logical_axes:
-        logical_axes = (None,) * ndim
-    elif len(logical_axes) < ndim:
-        logical_axes = logical_axes + (None,) * (ndim - len(logical_axes))
-    assert len(logical_axes) == ndim
+    assert logical_axes is not None, "Logical axes must be a tuple and cannot be None."
+    assert len(logical_axes) == ndim, "Logical axes must match the number of dimensions."
 
     non_contracting_dims = [i for i in range(ndim) if i not in contracting_dims]
     non_contracting_logical_axes = tuple(logical_axes[i] for i in non_contracting_dims)


Reviews:
user: login: sudhakarsingh27
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-13T23:11:44Z
user: login: jberchtold-nvidia
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-13T23:22:14Z
user: login: sudhakarsingh27
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-14T17:36:13Z
user: login: sudhakarsingh27
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-14T17:42:48Z
user: login: jberchtold-nvidia
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-14T18:43:16Z
user: login: jberchtold-nvidia
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-14T18:43:53Z
user: login: sudhakarsingh27
body: LGTM!

(Major change is replacing `param_with_axes` with `with_logical_partitioning` which enables MaxText runs with partitioning - especially TP)
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-05-16T17:32:58Z

Review comments:
diff_hunk: @@ -333,8 +333,8 @@ def get_non_contracting_logical_axes(ndim, logical_axes, contracting_dims):
         Tuple of logical axes for non-contracting dimensions.
path: transformer_engine/jax/sharding.py
user: login: sudhakarsingh27
body: This should be updated I guess then.
created_at: 2025-05-13T23:10:02Z
author_association: COLLABORATOR
diff_hunk: @@ -333,8 +333,8 @@ def get_non_contracting_logical_axes(ndim, logical_axes, contracting_dims):
         Tuple of logical axes for non-contracting dimensions.
     """
     if not logical_axes:
-        logical_axes = (None,) * ndim
-    elif len(logical_axes) < ndim:
path: transformer_engine/jax/sharding.py
user: login: sudhakarsingh27
body: Could you also explain the rationale here? 
created_at: 2025-05-13T23:10:16Z
author_association: COLLABORATOR
diff_hunk: @@ -333,8 +333,8 @@ def get_non_contracting_logical_axes(ndim, logical_axes, contracting_dims):
         Tuple of logical axes for non-contracting dimensions.
     """
     if not logical_axes:
-        logical_axes = (None,) * ndim
-    elif len(logical_axes) < ndim:
path: transformer_engine/jax/sharding.py
user: login: jberchtold-nvidia
body: We have two concepts here, and both use `None` so it's not the clearest. 

1. `logical_axes = None`, this is the default value for many of our Flax modules. For example, if `input_axes = None` as the default arg value in the constructor of the module, it means to not apply a sharding constraint to the input. At least that is my best guess on it's intention as I don't think it'd make sense to default to forcing tensors to be replicated. This is also supported by the fact that `with_sharding_constraint_by_logical_axes` is implemented as a no-op if the axes are `None`, meaning the existing sharding will apply and no additional constraints are applied.
2. `logical_axes = (None, None, ..., None)`, adding a sharding constraint to the tensor to make it replicated instead of partitioned. A tuple of `None`s per axis in the tensor's shape. This is distinct and not a default value anywhere and it means "enforce the tensor is replicated".

The function that I'm updating here, `get_non_contracting_logical_axes`, had originally converted (1) to (2). This means that if this function were called and the logical axes inputs were the default value, `None`, we would get out `(None, None, ..., None)` from this function. This means we are changing the behavior from (1) passing through a tensor without applying sharding constraints, the default, to (2) enforcing the sharding is replicated. This is not the intended behavior of this function.

This led to `ffn1`, the intermediate value inside the LayerNormMLP layer that is sharded according to outputs of `get_non_contracting_logical_axes`, to be replicated and led to us getting OOM on model configs that didn't with TE disabled.
created_at: 2025-05-13T23:22:14Z
author_association: COLLABORATOR
diff_hunk: @@ -333,8 +333,8 @@ def get_non_contracting_logical_axes(ndim, logical_axes, contracting_dims):
         Tuple of logical axes for non-contracting dimensions.
     """
     if not logical_axes:
-        logical_axes = (None,) * ndim
-    elif len(logical_axes) < ndim:
path: transformer_engine/jax/sharding.py
user: login: sudhakarsingh27
body: That's sneaky! Good catch
created_at: 2025-05-14T17:36:13Z
author_association: COLLABORATOR
diff_hunk: @@ -333,8 +333,8 @@ def get_non_contracting_logical_axes(ndim, logical_axes, contracting_dims):
         Tuple of logical axes for non-contracting dimensions.
     """
     if not logical_axes:
-        logical_axes = (None,) * ndim
-    elif len(logical_axes) < ndim:
path: transformer_engine/jax/sharding.py
user: login: sudhakarsingh27
body: Also, should we include this description in the code?
created_at: 2025-05-14T17:42:47Z
author_association: COLLABORATOR
diff_hunk: @@ -333,8 +333,8 @@ def get_non_contracting_logical_axes(ndim, logical_axes, contracting_dims):
         Tuple of logical axes for non-contracting dimensions.
path: transformer_engine/jax/sharding.py
user: login: jberchtold-nvidia
body: Switched this code to use an assert instead as none which better fit the use cases. So the doc comment is now correct
created_at: 2025-05-14T18:43:16Z
author_association: COLLABORATOR
diff_hunk: @@ -333,8 +333,8 @@ def get_non_contracting_logical_axes(ndim, logical_axes, contracting_dims):
         Tuple of logical axes for non-contracting dimensions.
     """
     if not logical_axes:
-        logical_axes = (None,) * ndim
-    elif len(logical_axes) < ndim:
path: transformer_engine/jax/sharding.py
user: login: jberchtold-nvidia
body: Updated `with_sharding_constraint_by_logical_axes` documentation to clarify this distinction
created_at: 2025-05-14T18:43:52Z
author_association: COLLABORATOR

PR comments:
user: login: jberchtold-nvidia
created_at: 2025-05-12T18:46:47Z
body: /te-ci L0
author_association: COLLABORATOR
user: login: jberchtold-nvidia
created_at: 2025-05-12T22:55:39Z
body: /te-ci L0
author_association: COLLABORATOR
user: login: jberchtold-nvidia
created_at: 2025-05-14T18:44:26Z
body: /te-ci L0
author_association: COLLABORATOR
user: login: jberchtold-nvidia
created_at: 2025-05-14T20:21:51Z
body: /te-ci L0
author_association: COLLABORATOR
user: login: jberchtold-nvidia
created_at: 2025-05-15T00:18:12Z
body: /te-ci L0
author_association: COLLABORATOR
user: login: jberchtold-nvidia
created_at: 2025-05-15T17:25:53Z
body: /te-ci L0
author_association: COLLABORATOR