nan in tests/jax/test_distributed_layernorm_mlp.py::TestDistributedLayernormMLP::test_layernorm_mlp_grad

Describe the bug

On one 8-GPU H100, with the ngc jax docker image nvcr.io/nvidia/jax:25.08-maxtext-py3, we were able to see nan in the single GPU forward value (the reference) in tests/jax/test_distributed_layernorm_mlp.py::TestDistributedLayernormMLP::test_layernorm_mlp_grad[with_jax_gemm_True-Float8CurrentScaling-use_bias_True-dtype_float16-activation_type_gelu-input_shape_4_128_256-fsdp1_tpsp2] using a pretty recent commit 7e593c3 (https://github.com/NVIDIA/TransformerEngine/tree/7e593c3be96b3eebc384da1a2ab307727065c9ab)

Steps/Code to reproduce bug
You can patch the following code diff and then run this exact pytest, it'll print single_fwd: nan

diff --git a/tests/jax/test_distributed_layernorm_mlp.py b/tests/jax/test_distributed_layernorm_mlp.py
index 667840da..0f6f77d0 100644
--- a/tests/jax/test_distributed_layernorm_mlp.py
+++ b/tests/jax/test_distributed_layernorm_mlp.py
@@ -220,82 +220,83 @@ class TestDistributedLayernormMLP:
                     static_argnums=range(len(inputs), len(static_inputs) + len(inputs)),
                 )
                 single_fwd, single_grads = single_jitter(*inputs, *static_inputs)
-
           # Multi GPUs
           devices = np.asarray(jax.devices()[:device_count]).reshape(*mesh_shape)
           mesh = Mesh(devices, mesh_axes)
           with mesh, autocast(
               enabled=quantization_recipe is not None,
               recipe=quantization_recipe,
               mesh_resource=mesh_resource,
           ):
               k1_sharding = NamedSharding(mesh, PartitionSpec("fsdp", None, "tpsp"))
               k2_sharding = NamedSharding(mesh, PartitionSpec("tpsp", "fsdp"))
               k1_ = jax.device_put(k1, k1_sharding)
               k2_ = jax.device_put(k2, k2_sharding)
               if use_bias:
                   b1_sharding = NamedSharding(mesh, PartitionSpec(None, "tpsp"))
                   b1_ = jax.device_put(b1, b1_sharding)
               else:
                   b1_sharding = b1_ = None
               multi_inputs = [*inputs[:2], k1_, k2_, b1_, *inputs[5:]]
-
               # Position ref for sharding pspec lists
               #   x, gamma, k1, k2, b1,
               #   b2
               in_shardings = (
                   None,
                   None,
                   k1_sharding,
                   k2_sharding,
                   b1_sharding,
                   None,
               )
               out_shardings = (
                   None,
                   (None, None, k1_sharding, k2_sharding, b1_sharding, None),
               )
-
               multi_jitter = jax.jit(
                   value_and_grad_func,
                   in_shardings=in_shardings,
                   out_shardings=out_shardings,
                   static_argnums=range(
                       len(multi_inputs), len(static_inputs) + len(multi_inputs) + 1
                   ),
               )  # +1 for multi_gpus
-
               multi_fwd, multi_grads = multi_jitter(*multi_inputs, *static_inputs, True)
-
       fwd_test_type = bwd_test_type = dtype
       if quantization_recipe is not None:
           quantize_config = get_quantize_config_with_recipe(quantization_recipe)
           fwd_test_type = quantize_config.FWD_DTYPE
           bwd_test_type = quantize_config.BWD_DTYPE
-
       if fwd_test_type == jnp.float16 and use_bias:
           assert_allclose(multi_fwd, single_fwd, atol=0.04, rtol=1.5)
       else:
           assert_allclose(multi_fwd, single_fwd, dtype=fwd_test_type)
-
       for i in range(len(inputs)):
           if multi_grads[i] is not None:
               if isinstance(multi_grads[i], list):
                   assert isinstance(single_grads[i], list)
                   for m_grad, s_grad in zip(multi_grads[i], single_grads[i]):
                       assert_allclose(
                           m_grad,
                           s_grad,
                           dtype=bwd_test_type,
                           err_msg=f"multi_grads[{i}] is not close",
                       )
               else:
                   assert_allclose(
                       multi_grads[i],
                       single_grads[i],
                       dtype=bwd_test_type,
                       err_msg=f"multi_grads[{i}] is not close",
                   )
+            jax.debug.print("single_fwd: {}", single_fwd)
+
+            ## Multi GPUs
+            #devices = np.asarray(jax.devices()[:device_count]).reshape(*mesh_shape)
+            #mesh = Mesh(devices, mesh_axes)
+            #with mesh, autocast(
+            #    enabled=quantization_recipe is not None,
+            #    recipe=quantization_recipe,
+            #    mesh_resource=mesh_resource,
+            #):
+            #    k1_sharding = NamedSharding(mesh, PartitionSpec("fsdp", None, "tpsp"))
+            #    k2_sharding = NamedSharding(mesh, PartitionSpec("tpsp", "fsdp"))
+            #    k1_ = jax.device_put(k1, k1_sharding)
+            #    k2_ = jax.device_put(k2, k2_sharding)
+            #    if use_bias:
+            #        b1_sharding = NamedSharding(mesh, PartitionSpec(None, "tpsp"))
+            #        b1_ = jax.device_put(b1, b1_sharding)
+            #    else:
+            #        b1_sharding = b1_ = None
+            #    multi_inputs = [*inputs[:2], k1_, k2_, b1_, *inputs[5:]]
+
+            #    # Position ref for sharding pspec lists
+            #    #   x, gamma, k1, k2, b1,
+            #    #   b2
+            #    in_shardings = (
+            #        None,
+            #        None,
+            #        k1_sharding,
+            #        k2_sharding,
+            #        b1_sharding,
+            #        None,
+            #    )
+            #    out_shardings = (
+            #        None,
+            #        (None, None, k1_sharding, k2_sharding, b1_sharding, None),
+            #    )
+
+            #    multi_jitter = jax.jit(
+            #        value_and_grad_func,
+            #        in_shardings=in_shardings,
+            #        out_shardings=out_shardings,
+            #        static_argnums=range(
+            #            len(multi_inputs), len(static_inputs) + len(multi_inputs) + 1
+            #        ),
+            #    )  # +1 for multi_gpus
+
+            #    multi_fwd, multi_grads = multi_jitter(*multi_inputs, *static_inputs, True)
+
+        #fwd_test_type = bwd_test_type = dtype
+        #if quantization_recipe is not None:
+        #    quantize_config = get_quantize_config_with_recipe(quantization_recipe)
+        #    fwd_test_type = quantize_config.FWD_DTYPE
+        #    bwd_test_type = quantize_config.BWD_DTYPE
+
+        #if fwd_test_type == jnp.float16 and use_bias:
+        #    assert_allclose(multi_fwd, single_fwd, atol=0.04, rtol=1.5)
+        #else:
+        #    assert_allclose(multi_fwd, single_fwd, dtype=fwd_test_type)
+
+        #for i in range(len(inputs)):
+        #    if multi_grads[i] is not None:
+        #        if isinstance(multi_grads[i], list):
+        #            assert isinstance(single_grads[i], list)
+        #            for m_grad, s_grad in zip(multi_grads[i], single_grads[i]):
+        #                assert_allclose(
+        #                    m_grad,
+        #                    s_grad,
+        #                    dtype=bwd_test_type,
+        #                    err_msg=f"multi_grads[{i}] is not close",
+        #                )
+        #        else:
+        #            assert_allclose(
+        #                multi_grads[i],
+        #                single_grads[i],
+        #                dtype=bwd_test_type,
+        #                err_msg=f"multi_grads[{i}] is not close",
+        #            )
 
     @pytest_parametrize_wrapper("mesh_config", generate_fsdp_and_tpsp_configs())
     @pytest_parametrize_wrapper("input_shape", INPUT_SHAPE)

Expected behavior
Not a nan fwd?

Environment overview (please complete the following information)

• Environment location: Docker with image nvcr.io/nvidia/jax:25.08-maxtext-py3
• Method of Transformer Engine install: pip install or from source with the following command: "pip install --no-build-isolation . -v"
• If method of install is [Docker], provide docker pull & docker run commands used:

Environment details

If NVIDIA docker image is used you don't need to specify these.
Otherwise, please provide:

• OS version
• PyTorch version
• Python version
• Transformer Engine version
• CUDA version
• CUDNN version

Device details

• GPU model

Additional context

Add any other context about the problem here.

Start with tests/jax/test_distributed_layernorm_mlp.py::TestDistributedLayernormMLP::test_layernorm_mlp_grad and reproduce the reported parameterized case in the nvcr.io/nvidia/jax:25.08-maxtext-py3 environment on an 8-GPU H100. Trace why the single-GPU forward reference becomes NaN; done means the test no longer produces an unexpected NaN and its forward and gradient comparisons pass.