Adding 2xACC control to FP8 GEMMs.

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

Author: mingxu1067

transformer_engine/jax/dot.py

@@ -11,6 +11,8 @@
 from .cpp_extensions import cast_transpose
 from .fp8 import FP8Helper, FP8MetaPackage
 
+Precision = jax.lax.Precision
+
 
 def type_safe_dot_general(
     x,
@@ -62,7 +64,8 @@ def fp8_dot_impl(
         lhs_scale_inv: jnp.ndarray,
         rhs_scale_inv: jnp.ndarray,
         ctype: jnp.dtype,    # computing type
-        contracting_dims: Tuple[Sequence[int], Sequence[int]]):
+        contracting_dims: Tuple[Sequence[int], Sequence[int]],
+        precision: Precision = None):
     """
     FP8 GEMM for XLA pattern match
     """
@@ -71,7 +74,14 @@ def fp8_dot_impl(
     lhs = dequantize(q_lhs, ctype, lhs_scale_inv)
     rhs = dequantize(q_rhs, ctype, rhs_scale_inv)
 
-    return jax.lax.dot_general(lhs, rhs, dim_nums)
+    return jax.lax.dot_general(lhs, rhs, dim_nums, precision=precision)
+
+
+def get_precision_of_fp8_dot(enable_2xACC: bool):
+    """
+    Get Precision of FP8 DOT.
+    """
+    return jax.lax.Precision.HIGHEST if enable_2xACC else jax.lax.Precision.DEFAULT
 
 
 @partial(jax.custom_vjp, nondiff_argnums=(6, 7, 8))
@@ -116,7 +126,8 @@ def _fp8_dot_fwd_rule(
     casted_kernel, updated_kernel_amax = quantize(kernel, fwd_dtype, kernel_scale)
 
     output = fp8_dot_impl(casted_x, casted_kernel, x_scale_inv, kernel_scale_inv, x.dtype,
-                          (lhs_contracting_dims, rhs_contracting_dims))
+                          (lhs_contracting_dims, rhs_contracting_dims),
+                          get_precision_of_fp8_dot(FP8Helper.FP8_2X_ACC_FPROP))
 
     ctx = (casted_x, casted_kernel, fp8_max, amax, scale, scale_inv, updated_x_amax,
            updated_kernel_amax, x.shape, kernel.shape)
@@ -144,14 +155,16 @@ def _fp8_dot_bwd_rule(fwd_dtype, bwd_dtype, contracting_dims, ctx, grad):    # p
     gt_constracting_dim = tuple(range(grad.ndim - len(x_constracting_dim), grad.ndim))
     x_scale_inv = scale_inv[gemm_x_idx]
     wgrad = fp8_dot_impl(casted_x, casted_grad_t, x_scale_inv, grad_scale_inv, grad.dtype,
-                         (x_constracting_dim, gt_constracting_dim))
+                         (x_constracting_dim, gt_constracting_dim),
+                         get_precision_of_fp8_dot(FP8Helper.FP8_2X_ACC_WGRAD))
 
     g_constracting_dim = tuple(
         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_kernel, grad_scale_inv, kernel_scale_inv, grad.dtype,
-                         (g_constracting_dim, k_constracting_dim))
+                         (g_constracting_dim, k_constracting_dim),
+                         get_precision_of_fp8_dot(FP8Helper.FP8_2X_ACC_DGRAD))
 
     amax = amax.at[gemm_x_idx, 0].set(updated_x_amax)
     amax = amax.at[gemm_kernel_idx, 0].set(updated_kernel_amax)

transformer_engine/jax/layernorm.py

@@ -10,7 +10,7 @@
 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
+from .dot import fp8_dot_impl, get_precision_of_fp8_dot
 from .fp8 import FP8Helper, FP8MetaPackage
 
 
@@ -193,7 +193,8 @@ def _layernorm_fp8_dot_fwd_rule(
 
     # (batch..., hidden_in) x (hidden_in, hidden_out...)
     output = fp8_dot_impl(ln_out, casted_kernel, x_scale_inv, kernel_scale_inv, x.dtype,
-                          (x_contracting_dims, k_contracting_dims))
+                          (x_contracting_dims, k_contracting_dims),
+                          get_precision_of_fp8_dot(FP8Helper.FP8_2X_ACC_FPROP))
 
     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,
@@ -231,14 +232,16 @@ def _layernorm_fp8_dot_bwd_rule(
     gt_constracting_dim = tuple(range(grad.ndim - len(xt_constracting_dim), grad.ndim))
     x_scale_inv = scale_inv[gemm_x_idx]
     wgrad = fp8_dot_impl(ln_out_t, casted_grad_t, x_scale_inv, grad_scale_inv, grad.dtype,
-                         (xt_constracting_dim, gt_constracting_dim))
+                         (xt_constracting_dim, gt_constracting_dim),
+                         get_precision_of_fp8_dot(FP8Helper.FP8_2X_ACC_WGRAD))
 
     g_for_dgrad_constracting_dim = tuple(
         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_kernel, grad_scale_inv, kernel_scale_inv, grad.dtype,
-                         (g_for_dgrad_constracting_dim, k_constracting_dim))
+                         (g_for_dgrad_constracting_dim, k_constracting_dim),
+                         get_precision_of_fp8_dot(FP8Helper.FP8_2X_ACC_DGRAD))
 
     if layernorm_type == 'layernorm':
         dx, dgamma, dbeta = layernorm_bwd(dgrad,

transformer_engine/jax/mlp.py

@@ -14,7 +14,7 @@
 from .cpp_extensions import dgated_gelu, dgated_gelu_cast_transpose
 from .cpp_extensions import rmsnorm_fwd_fp8, rmsnorm_bwd
 from .cpp_extensions import layernorm_fwd_fp8, layernorm_bwd
-from .dot import fp8_dot_impl, quantize
+from .dot import fp8_dot_impl, get_precision_of_fp8_dot, quantize
 from .layernorm import canonicalize_layernorm_type
 from .fp8 import FP8Helper, FP8MetaPackage
 
@@ -177,7 +177,8 @@ def _layernrom_geglu_fp8_mlp_fwd_rule(
 
     # (batch..., hidden_in) x (hidden_in, 2, hidden_out)
     dot_1_output = fp8_dot_impl(ln_out, casted_kernel_1, x_scale_inv, kernel_1_scale_inv, x.dtype,
-                                (x_contracting_dims, (0,)))
+                                (x_contracting_dims, (0,)),
+                                get_precision_of_fp8_dot(FP8Helper.FP8_2X_ACC_FPROP))
 
     gemm2_x_idx, gemm2_kernel_idx, _ = FP8Helper.get_fp8_meta_indices(1)
 
@@ -198,7 +199,8 @@ def _layernrom_geglu_fp8_mlp_fwd_rule(
 
     # (batch..., hidden_in) x (hidden_out, hidden_in)
     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,)))
+                                kernel_2_scale_inv, x.dtype, (x_contracting_dims, (0,)),
+                                get_precision_of_fp8_dot(FP8Helper.FP8_2X_ACC_FPROP))
 
     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,
@@ -237,12 +239,14 @@ def _layernrom_geglu_fp8_mlp_bwd_rule(
     # (hidden, batch...,) x (hidden, batch...)
     gemm2_x_scale_inv = scale_inv[gemm2_x_idx]
     wgrad_2 = fp8_dot_impl(casted_geglu_out_t, casted_grad_t, gemm2_x_scale_inv, grad_scale_inv,
-                           grad.dtype, (xt_batch_dims, xt_batch_dims))
+                           grad.dtype, (xt_batch_dims, xt_batch_dims),
+                           get_precision_of_fp8_dot(FP8Helper.FP8_2X_ACC_WGRAD))
 
     # (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_kernel_2, grad_scale_inv, kernel_2_scale_inv,
-                           grad.dtype, (x_contracting_dims, (1,)))
+                           grad.dtype, (x_contracting_dims, (1,)),
+                           get_precision_of_fp8_dot(FP8Helper.FP8_2X_ACC_DGRAD))
 
     gemm1_x_idx, gemm1_kernel_idx, gemm1_grad_idx = FP8Helper.get_fp8_meta_indices(0)
 
@@ -265,17 +269,16 @@ def _layernrom_geglu_fp8_mlp_bwd_rule(
     xt_batch_dims_plus_act_dim = tuple(i + 1 for i in xt_batch_dims)
     gemm1_x_scale_inv = scale_inv[gemm1_x_idx]
     wgrad_1 = fp8_dot_impl(ln_out_t, casted_dgeglu_t, gemm1_x_scale_inv, dgeglu_scale_inv,
-                           grad.dtype, (xt_batch_dims, xt_batch_dims_plus_act_dim))
+                           grad.dtype, (xt_batch_dims, xt_batch_dims_plus_act_dim),
+                           get_precision_of_fp8_dot(FP8Helper.FP8_2X_ACC_WGRAD))
 
     # (batch..., 2, hidden_out) x (hidden_in, 2, hidden_out)
     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_kernel_1, dgeglu_scale_inv, kernel_1_scale_inv,
-                           grad.dtype, (x_contracting_dims_plus_act_dim, (
-                               1,
-                               2,
-                           )))
+                           grad.dtype, (x_contracting_dims_plus_act_dim, (1, 2)),
+                           get_precision_of_fp8_dot(FP8Helper.FP8_2X_ACC_DGRAD))
 
     if layernorm_type == 'layernorm':
         dx, dgamma, dbeta = layernorm_bwd(dgrad_1,