[JAX] Refine MHA module

docs/api/jax.rst

@@ -32,6 +32,9 @@ Jax
 .. autoapiclass:: transformer_engine.jax.flax.RelativePositionBiases(num_buckets, max_distance, num_heads, **kwargs)
   :members: __call__
 
+.. autoapiclass:: transformer_engine.jax.flax.DotProductAttention(head_dim, num_heads, **kwargs)
+  :members: __call__
+
 .. autoapiclass:: transformer_engine.jax.flax.MultiHeadAttention(head_dim, num_heads, **kwargs)
   :members: __call__
 

tests/jax/test_custom_call_compute.py

@@ -14,13 +14,16 @@
 from flax import linen as nn
 
 from utils import assert_allclose
+from transformer_engine.jax.cpp_extensions import dgelu, dgelu_dbias_cast_transpose
+from transformer_engine.jax.cpp_extensions import gelu, gelu_fp8
 from transformer_engine.jax.cpp_extensions import dgated_gelu, gated_gelu
 from transformer_engine.jax.cpp_extensions import dgated_gelu_cast_transpose, gated_gelu_fp8
 from transformer_engine.jax.dot import type_safe_dot_general, dequantize, quantize
 from transformer_engine.jax.fp8 import FP8MetaPackage, FP8Helper
 from transformer_engine.jax.fp8 import is_fp8_available
 from transformer_engine.jax.layernorm import layernorm
 from transformer_engine.jax.mlp import layernrom_geglu_fp8_mlp
+from transformer_engine.jax.mlp import layernrom_gelu_fp8_mlp
 
 GEMM_CASES = [
     (256, 256, 512),
@@ -285,6 +288,126 @@ def ref_func(x, ln_s, y, z, fp8_max, fp8_metas_amax, fp8_metas_scale, fp8_metas_
                         jnp.asarray(ref_s_grad, np.float32),
                         dtype=FP8Helper.BWD_DTYPE)
 
+    @pytest.mark.skipif(not is_fp8_supported, reason=reason)
+    @pytest.mark.parametrize('m,n,k', [(256, 256, 512), (16384, 1024, 2816), (16384, 2816, 1024),
+                                       (16384, 1024, 1024)])
+    def test_grad_ln_gelu_fp8_mlp(self, m, n, k):
+        key = jax.random.PRNGKey(0)
+        subkeys = jax.random.split(key, 6)
+        activations = ('gelu',)
+
+        a = jax.random.normal(subkeys[0], (m, k), jnp.bfloat16)
+        k1 = jax.random.normal(subkeys[1], (k, len(activations), n), jnp.bfloat16)
+        k2 = jax.random.normal(subkeys[2], (n, k), jnp.bfloat16)
+        b1 = jax.random.normal(subkeys[3], (len(activations), n), jnp.bfloat16)
+        b2 = jax.random.normal(subkeys[4], (k,), jnp.bfloat16)
+        s = jax.random.normal(subkeys[5], (k,), jnp.bfloat16)
+
+        init_fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM * 2)
+        init_fp8_metas_amax = jnp.zeros(
+            (FP8Helper.NUM_META_PER_GEMM * 2, FP8Helper.AMAX_HISTORY_LEN), jnp.float32)
+        init_fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM * 2, 1), jnp.float32)
+        init_fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM * 2, 1), jnp.float32)
+
+        def primitive_func(x, ln_s, y, z, w, v, fp8_max, fp8_metas_amax, fp8_metas_scale,
+                           fp8_metas_scale_inv):
+            # x is input tensor, matrix 2d
+            # y, z are weights, matrix 2d
+            # out = ((x * y) + w) * z + v
+            fp8_meta_pkg = FP8MetaPackage(2, fp8_max, fp8_metas_amax, fp8_metas_scale,
+                                          fp8_metas_scale_inv)
+            return jnp.mean(
+                layernrom_gelu_fp8_mlp(x, ln_s, None, [y, z], [w, v], fp8_meta_pkg, "rmsnorm"))
+
+        def ln_gelu_fp8_mlp_ref(x: jnp.ndarray, ln_scale: jnp.ndarray, kernel_1: jnp.ndarray,
+                                kernel_2: jnp.ndarray, bias_1: jnp.ndarray, bias_2: jnp.ndarray,
+                                fp8_maxs: jnp.ndarray, amax: jnp.ndarray, scale: jnp.ndarray,
+                                scale_inv: jnp.ndarray) -> jnp.ndarray:
+
+            x = jnp.asarray(x, jnp.float32)
+            mean2 = jnp.mean(jax.lax.square(x), axis=-1, keepdims=True)
+            y = jnp.asarray(x * jax.lax.rsqrt(mean2 + 1e-6), jnp.bfloat16)
+            ln_out = y * ln_scale
+            ln_out = jnp.asarray(ln_out, jnp.bfloat16)
+
+            fp8_gemm_1_pkg = FP8MetaPackage(1, fp8_maxs[:FP8Helper.NUM_META_PER_GEMM],
+                                            amax[:FP8Helper.NUM_META_PER_GEMM],
+                                            scale[:FP8Helper.NUM_META_PER_GEMM],
+                                            scale_inv[:FP8Helper.NUM_META_PER_GEMM])
+            linear_1_out = type_safe_dot_general(ln_out, kernel_1, fp8_gemm_1_pkg, ((1,), (0,)))
+
+            bias_1_shape = (1,) * (linear_1_out.ndim - bias_1.ndim) + bias_1.shape
+            linear_1_out += jnp.reshape(bias_1, bias_1_shape)
+
+            x = jax.nn.gelu(linear_1_out)
+            x = jnp.asarray(jnp.squeeze(x, axis=-2), jnp.bfloat16)
+
+            fp8_gemm_2_pkg = FP8MetaPackage(1, fp8_maxs[FP8Helper.NUM_META_PER_GEMM:],
+                                            amax[FP8Helper.NUM_META_PER_GEMM:],
+                                            scale[FP8Helper.NUM_META_PER_GEMM:],
+                                            scale_inv[FP8Helper.NUM_META_PER_GEMM:])
+            output = type_safe_dot_general(x, kernel_2, fp8_gemm_2_pkg, ((1,), (0,)))
+
+            bias_2_shape = (1,) * (output.ndim - bias_2.ndim) + bias_2.shape
+            output += jnp.reshape(bias_2, bias_2_shape)
+
+            return output
+
+        def ref_func(x, ln_s, y, z, w, v, fp8_max, fp8_metas_amax, fp8_metas_scale,
+                     fp8_metas_scale_inv):
+            return jnp.mean(
+                ln_gelu_fp8_mlp_ref(x, ln_s, y, z, w, v, fp8_max, fp8_metas_amax, fp8_metas_scale,
+                                    fp8_metas_scale_inv))
+
+        value_n_grad_primitive_func = jit(
+            value_and_grad(primitive_func, (0, 1, 2, 3, 4, 5, 6, 7, 8, 9)))
+        value_n_grad_ref_func = jit(value_and_grad(ref_func, (0, 1, 2, 3, 4, 5, 6, 7, 8, 9)))
+
+        ref_fp8_max = init_fp8_max
+        ref_fp8_metas_amax = init_fp8_metas_amax
+        ref_fp8_metas_scale = init_fp8_metas_scale
+        ref_fp8_metas_scale_inv = init_fp8_metas_scale_inv
+
+        pri_fp8_max = init_fp8_max
+        pri_fp8_metas_amax = init_fp8_metas_amax
+        pri_fp8_metas_scale = init_fp8_metas_scale
+        pri_fp8_metas_scale_inv = init_fp8_metas_scale_inv
+
+        for _ in range(3):
+            ref_out, (ref_a_grad, ref_s_grad, ref_k1_grad, ref_k2_grad, ref_b1_grad, ref_b2_grad,
+                      ref_fp8_max, ref_fp8_metas_amax, ref_fp8_metas_scale,
+                      ref_fp8_metas_scale_inv) = value_n_grad_ref_func(
+                          a, s, k1, k2, b1, b2, ref_fp8_max, ref_fp8_metas_amax,
+                          ref_fp8_metas_scale, ref_fp8_metas_scale_inv)
+
+        for _ in range(3):
+            primitive_out, (primitive_a_grad, primitive_s_grad, primitive_k1_grad,
+                            primitive_k2_grad, primitive_b1_grad, primitive_b2_grad, pri_fp8_max,
+                            pri_fp8_metas_amax, pri_fp8_metas_scale,
+                            pri_fp8_metas_scale_inv) = value_n_grad_primitive_func(
+                                a, s, k1, k2, b1, b2, pri_fp8_max, pri_fp8_metas_amax,
+                                pri_fp8_metas_scale, pri_fp8_metas_scale_inv)
+
+        assert_allclose(primitive_out, ref_out, dtype=FP8Helper.FWD_DTYPE)
+        assert_allclose(jnp.asarray(primitive_a_grad, np.float32),
+                        jnp.asarray(ref_a_grad, np.float32),
+                        dtype=FP8Helper.BWD_DTYPE)
+        assert_allclose(jnp.asarray(primitive_k1_grad, np.float32),
+                        jnp.asarray(ref_k1_grad, np.float32),
+                        dtype=FP8Helper.BWD_DTYPE)
+        assert_allclose(jnp.asarray(primitive_k2_grad, np.float32),
+                        jnp.asarray(ref_k2_grad, np.float32),
+                        dtype=FP8Helper.BWD_DTYPE)
+        assert_allclose(jnp.asarray(primitive_s_grad, np.float32),
+                        jnp.asarray(ref_s_grad, np.float32),
+                        dtype=FP8Helper.BWD_DTYPE)
+        assert_allclose(jnp.asarray(primitive_b1_grad, np.float32),
+                        jnp.asarray(ref_b1_grad, np.float32),
+                        dtype=jnp.bfloat16)
+        assert_allclose(jnp.asarray(primitive_b2_grad, np.float32),
+                        jnp.asarray(ref_b2_grad, np.float32),
+                        dtype=jnp.bfloat16)
+
 
 @pytest.fixture(name="random_inputs")
 def random_inputs_fixture(shape):
@@ -294,6 +417,96 @@ def random_inputs_fixture(shape):
     return out
 
 
+class TestGeLu:
+
+    def ref_func(self, inputs):
+
+        func = jit(value_and_grad(lambda x: jnp.mean(jax.nn.gelu(x))))
+        return func(inputs)
+
+    def prim_func(self, inputs):
+
+        @jax.custom_vjp
+        def primitive(x):
+            out, _ = primitive_fwd(x)
+            return out
+
+        def primitive_fwd(x):
+            out = gelu(x)
+            ctx = x
+            return out, ctx
+
+        def primitive_bwd(ctx, g):
+            x = ctx
+            out = dgelu(g, x)
+            return (out,)
+
+        primitive.defvjp(primitive_fwd, primitive_bwd)
+        func = value_and_grad(lambda x: jnp.mean(primitive(x)))
+        return func(inputs)
+
+    @pytest.mark.parametrize('shape', [(32, 2, 64), (64, 2, 256)])
+    def test_gelu(self, random_inputs):
+        x = random_inputs
+        prim_out, prim_grad = self.prim_func(x)
+        ref_out, ref_grad = self.ref_func(x)
+
+        assert_allclose(prim_out, ref_out, dtype=x.dtype)
+        assert_allclose(prim_grad, ref_grad, dtype=x.dtype)
+
+
+class TestGeLuFP8(TestGeLu):
+
+    def prim_func(self, inputs):
+        amax = self.amax
+        scale = self.scale
+        scale_inv = self.scale_inv
+        no_use = jnp.zeros(1, jnp.float32)
+
+        @jax.custom_vjp
+        def primitive(x, y, z, w):
+            out = primitive_fwd(x)
+            return out
+
+        def primitive_fwd(x, y, z, w):
+            out, _ = gelu_fp8(x, amax, scale, scale_inv, jnp.float8_e4m3fn)
+            out = dequantize(out, x.dtype, scale_inv)
+            ctx = x
+            return out, ctx
+
+        def primitive_bwd(ctx, g):
+            x = ctx
+            dgelu, dgelu_trans, dbias, amax_out = dgelu_dbias_cast_transpose(
+                g, x, amax, scale, scale_inv, jnp.float8_e5m2, -1)
+            dgelu = dequantize(dgelu, x.dtype, scale_inv)
+            dgelu_trans = dequantize(dgelu_trans, x.dtype, scale_inv)
+            return dgelu, dgelu_trans, dbias, amax_out
+
+        primitive.defvjp(primitive_fwd, primitive_bwd)
+        func = value_and_grad(lambda x, y, z, w: jnp.mean(primitive(x, y, z, w)), (0, 1, 2, 3))
+
+        return func(inputs, no_use, no_use, no_use)
+
+    @pytest.mark.skipif(not is_fp8_supported, reason=reason)
+    @pytest.mark.parametrize('shape', [(32, 2, 64), (64, 2, 256)])
+    def test_gelu(self, random_inputs):
+        self.amax = jnp.zeros(1, jnp.float32)
+        self.scale = jnp.ones(1, jnp.float32)
+        self.scale_inv = jnp.ones(1, jnp.float32)
+
+        x = random_inputs
+        prim_out, (prim_grad, prim_grad_trans, dbias, amax) = self.prim_func(x)
+        ref_out, ref_grad = self.ref_func(x)
+
+        assert_allclose(prim_out, ref_out, dtype=FP8Helper.FWD_DTYPE)
+        assert_allclose(amax, jnp.amax(jnp.abs(ref_grad)), rtol=1e-2)
+        assert_allclose(dbias, jnp.sum(ref_grad, axis=(i for i in range(x.ndim - 1))))
+        assert_allclose(prim_grad, ref_grad, dtype=FP8Helper.BWD_DTYPE)
+        assert_allclose(prim_grad_trans,
+                        jnp.transpose(ref_grad, (2, 0, 1)),
+                        dtype=FP8Helper.BWD_DTYPE)
+
+
 class TestGatedGeLu:
 
     def ref_func(self, inputs):

tests/jax/test_layer.py

@@ -88,10 +88,12 @@ def compare_dict(ref_fd, test_fd, rtol=1e-05, atol=1e-08):
 _KEY_OF_SCALE_ATTN_LOGITS = "scale_attn_logits"
 _KEY_OF_NUM_HEADS = 'num_attention_heads'
 _KEY_OF_NUM_GQA_GROUPS = 'num_gqa_groups'
+_KEY_OF_ENABLE_ROPE = "enable_rotary_pos_emb"
 
 BASE_ATTRS = {
     _KEY_OF_TRANSPOSE_BS: True,
     _KEY_OF_NUM_HEADS: 8,
+    _KEY_OF_DROPOUT_RATE: 0,
 }
 
 ATTRS = [{
@@ -137,7 +139,25 @@ def compare_dict(ref_fd, test_fd, rtol=1e-05, atol=1e-08):
     _KEY_OF_FUSE_MLP_WI: True
 }, {
     _KEY_OF_NUM_HEADS: 8,
-    _KEY_OF_NUM_GQA_GROUPS: 4
+    _KEY_OF_NUM_GQA_GROUPS: 4,
+    _KEY_OF_TRANSPOSE_BS: False,
+    _KEY_OF_SCALE_ATTN_LOGITS: True,
+    _KEY_OF_LAYERNORM_TYPE: 'layernorm',
+    _KEY_OF_DROPOUT_RATE: 0.0,
+    _KEY_OF_MLP_ACTIVATIONS: (('gelu',)),
+    _KEY_OF_FUSE_MLP_WI: True
+}, {
+    _KEY_OF_TRANSPOSE_BS: False,
+    _KEY_OF_LAYERNORM_TYPE: 'layernorm',
+    _KEY_OF_DROPOUT_RATE: 0.0,
+    _KEY_OF_FUSE_MLP_WI: True,
+    _KEY_OF_ENABLE_ROPE: True
+}, {
+    _KEY_OF_TRANSPOSE_BS: True,
+    _KEY_OF_LAYERNORM_TYPE: 'layernorm',
+    _KEY_OF_DROPOUT_RATE: 0.0,
+    _KEY_OF_FUSE_MLP_WI: True,
+    _KEY_OF_ENABLE_ROPE: True
 }]
 
 ATTRS = [{**BASE_ATTRS, **attr} for attr in ATTRS]
@@ -202,7 +222,8 @@ def forward_runner(self, data_shape, dtype, attrs, rtol=1e-05, atol=1e-08):
         ref_out = loss_fn(inputs, ref_masks, ref_params, ref_others, ref_layer, apply_rng)
         test_out = loss_fn(inputs, test_masks, test_params, test_others, test_layer, apply_rng)
 
-        assert_allclose(ref_out, test_out, rtol=rtol, atol=atol)
+        if attrs[_KEY_OF_DROPOUT_RATE] == 0.:    # Skip elementwise checking for dropout
+            assert_allclose(ref_out, test_out, rtol=rtol, atol=atol)
 
         del data_rng, init_rng, apply_rng
 
@@ -263,9 +284,6 @@ def forward_backward_runner(self, data_shape, dtype, attrs, rtol=1e-05, atol=1e-
         test_out, test_grads = grad_fn(inputs, test_masks, test_params, test_others, test_layer,
                                        apply_rng)
 
-        assert_allclose(ref_out, test_out, rtol=rtol, atol=atol)
-        assert_allclose(ref_grads[0][0], test_grads[0][0], rtol=rtol, atol=atol)    # dgrad
-
         def reorganize_test_wgrad(test_wgrad, attrs):
             num_heads = attrs.get(_KEY_OF_NUM_HEADS)
             num_gqa_groups = attrs.get(_KEY_OF_NUM_GQA_GROUPS, num_heads)
@@ -309,10 +327,14 @@ def reorganize_test_wgrad(test_wgrad, attrs):
             del unfreeze_test_wgrad['mlp']['wo_kernel']
             return unfreeze_test_wgrad
 
-        compare_dict(ref_grads[1],
-                     reorganize_test_wgrad(test_grads[1], attrs),
-                     rtol=rtol,
-                     atol=atol)    # wgrad
+        if attrs[_KEY_OF_DROPOUT_RATE] == 0.:    # Skip elementwise checking for dropout
+            assert_allclose(ref_out, test_out, rtol=rtol, atol=atol)
+            assert_allclose(ref_grads[0][0], test_grads[0][0], rtol=rtol, atol=atol)    # dgrad
+
+            compare_dict(ref_grads[1],
+                         reorganize_test_wgrad(test_grads[1], attrs),
+                         rtol=rtol,
+                         atol=atol)    # wgrad
 
         del data_rng, init_rng, apply_rng
 
@@ -411,7 +433,8 @@ def forward_runner(self, data_shape, dtype, attrs, rtol=1e-05, atol=1e-08):
         ref_out = loss_fn(inputs, ref_masks, ref_params, ref_others, ref_layer, apply_rng)
         test_out = loss_fn(inputs, test_masks, test_params, test_others, test_layer, apply_rng)
 
-        assert_allclose(ref_out, test_out, rtol=rtol, atol=atol)
+        if attrs[_KEY_OF_DROPOUT_RATE] == 0.:    # Skip elementwise checking for dropout
+            assert_allclose(ref_out, test_out, rtol=rtol, atol=atol)
 
         del data_rng, init_rng, apply_rng
 
@@ -473,9 +496,6 @@ def forward_backward_runner(self, data_shape, dtype, attrs, rtol=1e-05, atol=1e-
         test_out, test_grads = grad_fn(inputs, test_masks, test_params, test_others, test_layer,
                                        apply_rng)
 
-        assert_allclose(ref_out, test_out, rtol=rtol, atol=atol)
-        assert_allclose(ref_grads[0][0], test_grads[0][0], rtol=rtol, atol=atol)    # dgrad
-
         def reorganize_test_wgrad(test_wgrad, attrs):
             num_heads = attrs.get(_KEY_OF_NUM_HEADS)
             num_gqa_groups = attrs.get(_KEY_OF_NUM_GQA_GROUPS, num_heads)
@@ -528,10 +548,13 @@ def reorganize_test_wgrad(test_wgrad, attrs):
             del unfreeze_test_wgrad['mlp']['wo_kernel']
             return unfreeze_test_wgrad
 
-        compare_dict(ref_grads[1],
-                     reorganize_test_wgrad(test_grads[1], attrs),
-                     rtol=rtol,
-                     atol=atol)    # wgrad
+        if attrs[_KEY_OF_DROPOUT_RATE] == 0.:    # Skip elementwise checking for dropout
+            assert_allclose(ref_out, test_out, rtol=rtol, atol=atol)
+            assert_allclose(ref_grads[0][0], test_grads[0][0], rtol=rtol, atol=atol)    # dgrad
+            compare_dict(ref_grads[1],
+                         reorganize_test_wgrad(test_grads[1], attrs),
+                         rtol=rtol,
+                         atol=atol)    # wgrad
 
         del data_rng, init_rng, apply_rng