Modifiied code with the review feedback.

Signed-off-by: Ming Huang <[email protected]>

Author: mingxu1067

transformer_engine/jax/cpp_extensions.py

@@ -2884,7 +2884,7 @@ def cast_transpose(x: jnp.ndarray, amax: jnp.ndarray, scale: jnp.ndarray, scale_
         transpose_axis_boundary=transpose_axis_boundary)
 
 
-class CastPrimitive(BasePrimitive):
+class CastFP8Primitive(BasePrimitive):
     """
     Cast Primitive
     """
@@ -2941,7 +2941,11 @@ def lowering(ctx, x, amax, scale, scale_inv, *, out_dtype):
                                                                jax_dtype_to_te_dtype(x_aval.dtype),
                                                                jax_dtype_to_te_dtype(out_dtype))
 
-        out = custom_caller(CastPrimitive.name, args, opaque, False, operand_output_aliases={1: 1})
+        out = custom_caller(CastFP8Primitive.name,
+                            args,
+                            opaque,
+                            False,
+                            operand_output_aliases={1: 1})
 
         return out
 
@@ -2950,23 +2954,23 @@ def impl(x, amax, scale, scale_inv, out_dtype):
         """
         te_cast implementation
         """
-        assert CastPrimitive.inner_primitive is not None
+        assert CastFP8Primitive.inner_primitive is not None
         casted_x, updated_amax = \
-            CastPrimitive.inner_primitive.bind(
+            CastFP8Primitive.inner_primitive.bind(
                 x, amax, scale, scale_inv, out_dtype=out_dtype)
         return casted_x, updated_amax
 
     @staticmethod
     def batcher(batched_args, batch_dims, *, out_dtype):
         _check_valid_batch_dims(batch_dims)
-        assert CastPrimitive.outer_primitive is not None
+        assert CastFP8Primitive.outer_primitive is not None
 
         x, amax, scale, scale_inv = batched_args
         x_bdim, amax_bdim, *_ = batch_dims
 
         out_bdims = x_bdim, x_bdim, amax_bdim
-        return CastPrimitive.outer_primitive.bind(x, amax, scale, scale_inv,
-                                                  out_dtype=out_dtype), out_bdims
+        return CastFP8Primitive.outer_primitive.bind(x, amax, scale, scale_inv,
+                                                     out_dtype=out_dtype), out_bdims
 
     @staticmethod
     def infer_sharding_from_operands(out_dtype, mesh, arg_infos, result_infos):
@@ -2987,24 +2991,24 @@ def partition(out_dtype, mesh, arg_infos, result_infos):
 
         def sharded_impl(x, amax, scale, scale_inv):
             local_cx, local_updated_amax = \
-                CastPrimitive.impl(x, amax, scale, scale_inv, out_dtype=out_dtype)
+                CastFP8Primitive.impl(x, amax, scale, scale_inv, out_dtype=out_dtype)
             global_updated_amax = all_reduce_max_along_all_axes_except_PP(local_updated_amax)
 
             return local_cx, global_updated_amax
 
         return mesh, sharded_impl, out_shardings, arg_shardings
 
 
-register_primitive(CastPrimitive)
+register_primitive(CastFP8Primitive)
 
 
-def cast(x: jnp.ndarray, amax: jnp.ndarray, scale: jnp.ndarray, scale_inv: jnp.ndarray,
-         out_dtype: TEDType) -> Tuple[jnp.ndarray, jnp.ndarray]:
+def cast_fp8(x: jnp.ndarray, amax: jnp.ndarray, scale: jnp.ndarray, scale_inv: jnp.ndarray,
+             out_dtype: TEDType) -> Tuple[jnp.ndarray, jnp.ndarray]:
     """
     Cast wrapper
     Return FP8 tensor
     """
-    return CastPrimitive.outer_primitive.bind(x, amax, scale, scale_inv, out_dtype=out_dtype)
+    return CastFP8Primitive.outer_primitive.bind(x, amax, scale, scale_inv, out_dtype=out_dtype)
 
 
 class TransposePrimitive(BasePrimitive):

transformer_engine/jax/dot.py

@@ -113,20 +113,20 @@ def _fp8_dot_fwd_rule(
     kernel_scale_inv = scale_inv[gemm_kernel_idx]
     # Note (Ming Huang): Use native cast to allow XLA handle tranpose for avoiding
     # unnecessary copy to break FP8 GEMM pattern matching.
-    casted_kerenl, updated_kernel_amax = quantize(kernel, fwd_dtype, kernel_scale)
+    casted_kernel, updated_kernel_amax = quantize(kernel, fwd_dtype, kernel_scale)
 
-    output = fp8_dot_impl(casted_x, casted_kerenl, x_scale_inv, kernel_scale_inv, x.dtype,
+    output = fp8_dot_impl(casted_x, casted_kernel, x_scale_inv, kernel_scale_inv, x.dtype,
                           (lhs_contracting_dims, rhs_contracting_dims))
 
-    ctx = (casted_x, casted_kerenl, fp8_max, amax, scale, scale_inv, updated_x_amax,
+    ctx = (casted_x, casted_kernel, fp8_max, amax, scale, scale_inv, updated_x_amax,
            updated_kernel_amax, x.shape, kernel.shape)
     return output, ctx
 
 
 def _fp8_dot_bwd_rule(fwd_dtype, bwd_dtype, contracting_dims, ctx, grad):    # pylint: disable=unused-argument
     lhs_contracting_dims, rhs_contracting_dims = contracting_dims
 
-    casted_x, casted_kerenl, fp8_max, amax, scale, scale_inv, \
+    casted_x, casted_kernel, fp8_max, amax, scale, scale_inv, \
         updated_x_amax, updated_kernel_amax, x_shape, kernel_shape = ctx
 
     gemm_x_idx, gemm_kernel_idx, gemm_grad_idx = FP8Helper.get_fp8_meta_indices(0)
@@ -150,7 +150,7 @@ def _fp8_dot_bwd_rule(fwd_dtype, bwd_dtype, contracting_dims, ctx, grad):    # p
         range(grad.ndim - len(kernel_shape) + len(rhs_contracting_dims), grad.ndim))
     k_constracting_dim = tuple(range(len(rhs_contracting_dims), len(kernel_shape)))
     kernel_scale_inv = scale_inv[gemm_kernel_idx]
-    dgrad = fp8_dot_impl(casted_grad, casted_kerenl, grad_scale_inv, kernel_scale_inv, grad.dtype,
+    dgrad = fp8_dot_impl(casted_grad, casted_kernel, grad_scale_inv, kernel_scale_inv, grad.dtype,
                          (g_constracting_dim, k_constracting_dim))
 
     amax = amax.at[gemm_x_idx, 0].set(updated_x_amax)

transformer_engine/jax/layernorm.py

@@ -7,7 +7,7 @@
 import jax
 import jax.numpy as jnp
 
-from .cpp_extensions import cast, cast_transpose, transpose
+from .cpp_extensions import cast_fp8, cast_transpose, transpose
 from .cpp_extensions import rmsnorm_fwd, rmsnorm_fwd_fp8, rmsnorm_bwd
 from .cpp_extensions import layernorm_fwd, layernorm_fwd_fp8, layernorm_bwd
 from .dot import fp8_dot_impl
@@ -188,14 +188,14 @@ def _layernorm_fp8_dot_fwd_rule(
     # Kernel in (hidden_in, hidden_out...)
     # Note (Ming Huang): Use cast only to allow XLA handle tranpose for avoiding
     # unnecessary copy to break FP8 GEMM pattern matching.
-    casted_kerenl, updated_kernel_amax = \
-        cast(kernel, kernel_amax, kernel_scale, kernel_scale_inv, fwd_dtype)
+    casted_kernel, updated_kernel_amax = \
+        cast_fp8(kernel, kernel_amax, kernel_scale, kernel_scale_inv, fwd_dtype)
 
     # (batch..., hidden_in) x (hidden_in, hidden_out...)
-    output = fp8_dot_impl(ln_out, casted_kerenl, x_scale_inv, kernel_scale_inv, x.dtype,
+    output = fp8_dot_impl(ln_out, casted_kernel, x_scale_inv, kernel_scale_inv, x.dtype,
                           (x_contracting_dims, k_contracting_dims))
 
-    ctx = (ln_out, casted_kerenl, fp8_max, amax, scale, scale_inv, updated_x_amax,
+    ctx = (ln_out, casted_kernel, fp8_max, amax, scale, scale_inv, updated_x_amax,
            updated_kernel_amax, x.shape, kernel.shape, mu, rsigma, x, gamma, x_contracting_dims,
            k_contracting_dims)
 
@@ -210,7 +210,7 @@ def _layernorm_fp8_dot_bwd_rule(
         epsilon,
         ctx,
         grad):
-    ln_out_, casted_kerenl, fp8_max, amax, scale, scale_inv, \
+    ln_out_, casted_kernel, fp8_max, amax, scale, scale_inv, \
     updated_x_amax, updated_kernel_amax, \
     x_shape, kernel_shape, mu, rsigma, x, gamma, \
     x_contracting_dims, k_contracting_dims = ctx
@@ -237,7 +237,7 @@ def _layernorm_fp8_dot_bwd_rule(
         range(grad.ndim - len(kernel_shape) + len(k_contracting_dims), grad.ndim))
     k_constracting_dim = tuple(range(len(k_contracting_dims), len(kernel_shape)))
     kernel_scale_inv = scale_inv[gemm_kernel_idx]
-    dgrad = fp8_dot_impl(casted_grad, casted_kerenl, grad_scale_inv, kernel_scale_inv, grad.dtype,
+    dgrad = fp8_dot_impl(casted_grad, casted_kernel, grad_scale_inv, kernel_scale_inv, grad.dtype,
                          (g_for_dgrad_constracting_dim, k_constracting_dim))
 
     if layernorm_type == 'layernorm':

transformer_engine/jax/mlp.py

@@ -9,7 +9,7 @@
 import jax
 import jax.numpy as jnp
 
-from .cpp_extensions import cast, transpose, cast_transpose
+from .cpp_extensions import cast_fp8, transpose, cast_transpose
 from .cpp_extensions import gated_gelu, gated_gelu_fp8
 from .cpp_extensions import dgated_gelu, dgated_gelu_cast_transpose
 from .cpp_extensions import rmsnorm_fwd_fp8, rmsnorm_bwd
@@ -172,11 +172,11 @@ def _layernrom_geglu_fp8_mlp_fwd_rule(
 
     # Note (Ming Huang): Use cast only to allow XLA handle tranpose for avoiding
     # unnecessary copy to break FP8 GEMM pattern matching.
-    casted_kerenl_1, updated_kernel_1_amax = \
-        cast(kernel_1, kernel_1_amax, kernel_1_scale, kernel_1_scale_inv, fwd_dtype)
+    casted_kernel_1, updated_kernel_1_amax = \
+        cast_fp8(kernel_1, kernel_1_amax, kernel_1_scale, kernel_1_scale_inv, fwd_dtype)
 
     # (batch..., hidden_in) x (hidden_in, 2, hidden_out)
-    dot_1_output = fp8_dot_impl(ln_out, casted_kerenl_1, x_scale_inv, kernel_1_scale_inv, x.dtype,
+    dot_1_output = fp8_dot_impl(ln_out, casted_kernel_1, x_scale_inv, kernel_1_scale_inv, x.dtype,
                                 (x_contracting_dims, (0,)))
 
     gemm2_x_idx, gemm2_kernel_idx, _ = FP8Helper.get_fp8_meta_indices(1)
@@ -194,14 +194,14 @@ def _layernrom_geglu_fp8_mlp_fwd_rule(
     kernel_2_scale_inv = scale_inv[gemm2_kernel_idx]
     # Note (Ming Huang): Use native cast to allow XLA handle tranpose for avoiding
     # unnecessary copy to break FP8 GEMM pattern matching.
-    casted_kerenl_2, updated_kernel_2_amax = quantize(kernel_2, fwd_dtype, kernel_2_scale)
+    casted_kernel_2, updated_kernel_2_amax = quantize(kernel_2, fwd_dtype, kernel_2_scale)
 
     # (batch..., hidden_in) x (hidden_out, hidden_in)
-    dot_2_output = fp8_dot_impl(casted_geglu_out, casted_kerenl_2, geglu_out_scale_inv,
+    dot_2_output = fp8_dot_impl(casted_geglu_out, casted_kernel_2, geglu_out_scale_inv,
                                 kernel_2_scale_inv, x.dtype, (x_contracting_dims, (0,)))
 
-    ctx = (x, ln_out, mu, rsigma, gamma, dot_1_output, casted_geglu_out, casted_kerenl_1,
-           casted_kerenl_2, fp8_max, amax, scale, scale_inv, updated_x_amax, updated_geglu_amax,
+    ctx = (x, ln_out, mu, rsigma, gamma, dot_1_output, casted_geglu_out, casted_kernel_1,
+           casted_kernel_2, fp8_max, amax, scale, scale_inv, updated_x_amax, updated_geglu_amax,
            updated_kernel_1_amax, updated_kernel_2_amax, x_contracting_dims, xt_batch_dims)
 
     return dot_2_output, ctx
@@ -216,7 +216,7 @@ def _layernrom_geglu_fp8_mlp_bwd_rule(
         ctx,
         grad):
     x, ln_out, mu, rsigma, gamma, dot_1_output, casted_geglu_out, \
-    casted_kerenl_1, casted_kerenl_2, fp8_max, amax, scale, scale_inv, updated_x_amax, \
+    casted_kernel_1, casted_kernel_2, fp8_max, amax, scale, scale_inv, updated_x_amax, \
     updated_geglu_amax, updated_kernel_1_amax, updated_kernel_2_amax, \
     x_contracting_dims, xt_batch_dims = ctx
 
@@ -241,7 +241,7 @@ def _layernrom_geglu_fp8_mlp_bwd_rule(
 
     # (batch..., hidden_out) x (hidden_in, hidden_out)
     kernel_2_scale_inv = scale_inv[gemm2_kernel_idx]
-    dgrad_2 = fp8_dot_impl(casted_grad, casted_kerenl_2, grad_scale_inv, kernel_2_scale_inv,
+    dgrad_2 = fp8_dot_impl(casted_grad, casted_kernel_2, grad_scale_inv, kernel_2_scale_inv,
                            grad.dtype, (x_contracting_dims, (1,)))
 
     gemm1_x_idx, gemm1_kernel_idx, gemm1_grad_idx = FP8Helper.get_fp8_meta_indices(0)
@@ -271,7 +271,7 @@ def _layernrom_geglu_fp8_mlp_bwd_rule(
     x_contracting_dims_plus_act_dim = (min(x_contracting_dims),) + tuple(
         i + 1 for i in x_contracting_dims)
     kernel_1_scale_inv = scale_inv[gemm1_kernel_idx]
-    dgrad_1 = fp8_dot_impl(casted_dgeglu, casted_kerenl_1, dgeglu_scale_inv, kernel_1_scale_inv,
+    dgrad_1 = fp8_dot_impl(casted_dgeglu, casted_kernel_1, dgeglu_scale_inv, kernel_1_scale_inv,
                            grad.dtype, (x_contracting_dims_plus_act_dim, (
                                1,
                                2,