[JAX] Regression tests for custom ops with jax.experimental.custom_partitioning (#471)

[JAX] Regression tests for custom ops sharding with both xmap and custom_partitioning.

Coverage:
- layernorm: fwd/grad, zero_centered_gamma, DP, TP_COL, DP_TP_COL
- rmsnorm: fwd/grad, DP, TP_COL, DP_TP_COL
- softmax: fwd/grad, SCALED, SCALED_MASKED, SCALED_UPPER_TRIANG_MASKED, DP, TP_COL, TP_ROW, DP_TP_COL, DP_TP_ROW
- self_fused_attn: fwd/grad, NO_BIAS, PRE_SCALE_BIAS, POST_SCALE_BIAS, NO_MASK, CAUSAL_MASK, PADDING_MASK, DP, TP_COL, DP_TP_COL
- cross_fused_attn: fwd/grad, NO_BIAS, NO_MASK, PADDING_MASK, DP, TP_COL, DP_TP_COL
Signed-off-by: Alp Dener <adener@nvidia.com>

qa/L0_jax_distributed_unittest/test.sh

+# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+set -xe
+
+: ${TE_PATH:=/opt/transformerengine}
+pytest -Wignore -v $TE_PATH/tests/jax/test_distributed_custom_ops.py
+

tests/jax/custom_ops_helper.py

+# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+import os
+import pytest
+import numpy as np
+from dataclasses import dataclass
+from typing import Tuple
+from enum import Enum
+from functools import partial
+
+import jax
+import jax.numpy as jnp
+from jax import random
+from jax.experimental.pjit import pjit, _UNSPECIFIED
+from jax.sharding import PartitionSpec
+
+import flax
+
+from transformer_engine.jax.sharding import ShardingType
+try:
+    # try importing the new custom partitioning implementation
+    from transformer_engine.jax.sharding import MeshResource
+except ImportError:
+    # must be using an older TE/JAX version so fall back on the xmap sharding implementation
+    MeshResource = None
+from transformer_engine.jax.layernorm import layernorm
+from transformer_engine.jax.softmax import SoftmaxType, softmax
+from transformer_engine.jax.fused_attn import \
+    AttnBiasType, AttnMaskType, is_fused_attn_kernel_available, self_fused_attn, cross_fused_attn
+
+
+class FusedAttnBackend(Enum):
+        Max512 = "0"
+        Arbitrary = "1"
+
+
+@pytest.fixture(name="backend", params=[FusedAttnBackend.Max512, FusedAttnBackend.Arbitrary])
+def fixture_backend(request):
+    backend = request.param
+    os.environ["NVTE_FUSED_ATTN_BACKEND"] = backend.value
+    yield backend
+    os.environ["NVTE_FUSED_ATTN_BACKEND"] = ""
+
+
+@dataclass
+class CustomOpsTestHelper:
+    qkv_shape: Tuple[int,int,int,int] = (32, 128, 16, 64)
+    pad_ratio: float = 0.3
+    dropout_prob: float = 0.1
+    dtype: type = jnp.float16
+    
+    @staticmethod
+    def use_custom_partitioning():
+        return (MeshResource is not None)
+
+    @staticmethod
+    def get_sharding_spec(mesh_names, sharding_type):
+        P = PartitionSpec
+        if sharding_type is ShardingType.DP:
+            return P(mesh_names[0], None), P(None), P(None)
+        elif sharding_type is ShardingType.DP_TP_COL:
+            return P(mesh_names[0], mesh_names[1]), P(None), P(None)
+        else:
+            raise NotImplementedError
+    
+    @staticmethod
+    def get_sharding_resource(mesh_names, sharding_type):
+        dp_r = None
+        tp_r = None
+        
+        if sharding_type in (ShardingType.DP, ShardingType.DP_TP_COL, ShardingType.DP_TP_ROW):
+            dp_r = mesh_names[0]
+        if sharding_type in (ShardingType.TP_COL, ShardingType.TP_ROW):
+            tp_r = mesh_names[0]
+        if sharding_type in (ShardingType.DP_TP_COL, ShardingType.DP_TP_ROW):
+            tp_r = mesh_names[1]
+
+        return MeshResource(dp_r, tp_r)
+
+    @staticmethod
+    def make_mask(q_tokens, kv_tokens, mask_type, dtype=jnp.uint8):
+        if mask_type == AttnMaskType.CAUSAL_MASK:
+            causal = flax.linen.make_causal_mask(q_tokens, dtype=dtype)
+            padding = flax.linen.make_attention_mask(q_tokens > 0, kv_tokens > 0, dtype=dtype)
+            return flax.linen.combine_masks(causal, padding)
+        else:
+            return flax.linen.make_attention_mask(q_tokens > 0, kv_tokens > 0, dtype=dtype)
+
+    @staticmethod
+    def count_collectives(hlo):
+        tmp = hlo.splitlines()
+        symb = "-start"
+        result = {
+            "all-reduce" : 0,
+            "other" : 0
+        }
+        for line in tmp:
+            txt = line.split()
+            if len(txt) > 0 and symb in txt[0]:
+                if "all-reduce" in txt[0]:
+                    result["all-reduce"] += 1
+                else:
+                    result["other"] += 1
+        return result
+
+    def get_tolerance(self, ref_val, relaxation=2./3., dtype=None):
+        if dtype is None:
+            dtype = self.dtype
+
+        # slightly relax the machine epsilon for minimum tolerance
+        eps_relaxed = jax.lax.pow(jnp.finfo(dtype).eps, dtype(relaxation))
+
+        # calculate the "Unit of Least Precision" -- i.e. distance to the next representable number
+        spacing_high = jnp.nextafter(dtype(ref_val), jnp.finfo(dtype).max) - dtype(ref_val)
+        spacing_low = dtype(ref_val) - jnp.nextafter(dtype(ref_val, jnp.finfo(dtype).min))
+        ulp = jax.lax.max(spacing_low, spacing_high)
+
+        return jax.lax.max(eps_relaxed, ulp)
+
+    def compare_ops(self, custom_func, ref_func, ref_count,
+                    *args, grad_args=None, dtype=None,
+                    in_shardings=_UNSPECIFIED, out_shardings=_UNSPECIFIED,
+                    **kwargs):
+        if dtype is None:
+            dtype = self.dtype
+
+        if isinstance(custom_func, partial):
+            func_name = custom_func.func.__name__
+        else:
+            func_name = custom_func.__name__
+        func_name = func_name.removeprefix('custom_')
+        if grad_args is None:
+            grad_args = tuple(range(len(args)))
+        
+        custom_gradded = jax.value_and_grad(custom_func, argnums=grad_args)
+        test_fwd, test_grads = custom_gradded(*args, **kwargs)
+        custom_pjitter = pjit(custom_gradded,
+                              in_shardings=in_shardings,
+                              out_shardings=out_shardings)
+        custom_hlo = custom_pjitter.lower(*args, **kwargs).compile().as_text()
+        custom_count = self.count_collectives(custom_hlo)
+        if ref_count is not None:
+            assert custom_count==ref_count, \
+                f"`{func_name}`: Expected collective count is {ref_count}, but got {custom_count}."
+        else:
+            print(f"`{func_name}`: Output collective count is {custom_count}.")
+
+        ref_gradded = jax.value_and_grad(ref_func, argnums=grad_args)
+        ref_fwd, ref_grads = ref_gradded(*args, **kwargs)
+        fwd_tol = self.get_tolerance(ref_fwd, dtype=dtype)
+        assert jnp.allclose(test_fwd, ref_fwd, rtol=0.0, atol=fwd_tol), \
+            f"`{func_name}`: Output (fwd) error {jnp.max(jnp.abs(test_fwd - ref_fwd))}" + \
+            f" exceeds tolerance ({fwd_tol})."
+
+        if len(grad_args) == 1:
+            ref_grads = (ref_grads, )
+            test_grads = (test_grads, )
+        failed_grads = {}
+        for i, grads in enumerate(zip(test_grads, ref_grads)):
+            test_grad, ref_grad = grads
+            if test_grad is None and ref_grad is None:
+                continue
+            bwd_tol = self.get_tolerance(ref_grad, dtype=dtype)
+            if not jnp.allclose(test_grad, ref_grad, rtol=0.0, atol=bwd_tol):
+                failed_grads[i] = jnp.max(jnp.abs(test_grad - ref_grad))
+        assert len(failed_grads) == 0, \
+            f"`{func_name}`: Gradient (bwd) max errors" + \
+            f" [{', '.join([f'Arg{k}={v}' for k,v in failed_grads.items()])}]" + \
+            f" exceed tolerance ({bwd_tol})."
+
+    @staticmethod
+    def check_fused_attn_inputs(self, q_seq, kv_seq, head_dim, pad_ratio, dropout_probability,
+                                attn_bias_type, attn_mask_type, backend, dtype=jnp.float16):
+        if (q_seq > 512 or kv_seq > 512 or backend == FusedAttnBackend.Arbitrary) \
+            and pad_ratio != 0:
+            pytest.skip(
+                "`fused_attention`: Arbitrary seqlen backend does not support padded input.")
+
+        if not is_fused_attn_kernel_available(
+            dtype, dtype, attn_bias_type, attn_mask_type,
+            dropout_probability, q_seq, kv_seq, head_dim):
+            pytest.skip(
+                "`fused_attention`: Unsupported inputs combination or device compute capability.")
+
+    def fused_attn_core(self, query, key, value, bias, mask, scale_factor,
+                        attn_bias_type, attn_mask_type, dropout_rng, dropout_prob):
+        # Q*K matmul
+        query = jnp.squeeze(query)
+        key = jnp.squeeze(key)
+        value = jnp.squeeze(value)
+        attn_weights = jnp.einsum("...qhd,...khd->...hqk", query, key)
+        # scale and bias
+        if attn_bias_type == AttnBiasType.PRE_SCALE_BIAS:
+            attn_weights = scale_factor * (attn_weights + bias)
+        elif attn_bias_type == AttnBiasType.POST_SCALE_BIAS:
+            attn_weights = scale_factor * attn_weights + bias
+        else:
+            attn_weights = scale_factor * attn_weights
+        # padding mask
+        if attn_mask_type != AttnMaskType.NO_MASK and mask is not None:
+            big_neg = jnp.finfo(query.dtype).min
+            attn_weights = jnp.where(mask, attn_weights, big_neg)
+        # softmax
+        attn_weights = jax.nn.softmax(attn_weights).astype(query.dtype)
+        # dropout
+        if dropout_prob == 1.0:
+            attn_weights = jnp.zeros_like(attn_weights)
+        elif dropout_prob > 0.0:
+            keep_prob = 1.0 - dropout_prob
+            keep = random.bernoulli(dropout_rng, p=keep_prob, shape=attn_weights.shape)
+            multiplier = keep.astype(query.dtype) / jnp.asarray(keep_prob, dtype=query.dtype)
+            attn_weights = attn_weights * multiplier
+        # QK*V matmul
+        result =  jnp.einsum('...hqk,...khd->...qhd', attn_weights, value)
+        return jnp.mean(result)
+
+    @staticmethod
+    def custom_layernorm(x, gamma, beta, zero_centered_gamma, epsilon, sharding_type):
+        result = layernorm(x, gamma, beta,
+                           layernorm_type='layernorm',
+                           zero_centered_gamma=zero_centered_gamma,
+                           epsilon=epsilon,
+                           sharding_type=sharding_type,
+                           dp_dim_index=0)
+        return jnp.mean(result)
+
+    def reference_layernorm(self, x, gamma, beta, zero_centered_gamma, epsilon):
+        x_ = jnp.asarray(x, jnp.float32)
+        mean = jnp.mean(x_, axis=-1, keepdims=True)
+        var = jnp.mean(jnp.square(x_ - mean), axis=-1, keepdims=True)
+        normed_input = (x_ - mean) * jax.lax.rsqrt(var + epsilon)
+        if zero_centered_gamma:
+            result = jnp.asarray(normed_input * (gamma + 1) + beta).astype(x.dtype)
+        else:
+            result = jnp.asarray(normed_input * gamma + beta).astype(x.dtype)
+        return jnp.mean(result)
+
+    @staticmethod
+    def custom_rmsnorm(x, gamma, epsilon, sharding_type):
+        result = layernorm(x, gamma, None,
+                           layernorm_type='rmsnorm',
+                           zero_centered_gamma=False,
+                           epsilon=epsilon,
+                           sharding_type=sharding_type,
+                           dp_dim_index=0)
+        return jnp.mean(result)
+
+    def reference_rmsnorm(self, x, gamma, epsilon):
+        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 + epsilon), x.dtype)
+        result = y * gamma
+        return jnp.mean(result)
+
+    @staticmethod
+    def custom_softmax(x, mask, scale_factor, softmax_type, sharding_type):
+        result = softmax(x, mask,
+                         scale_factor=scale_factor, 
+                         softmax_type=softmax_type,
+                         sharding_type=sharding_type)
+        return jnp.mean(result)
+
+    def reference_softmax(self, x, mask, scale_factor, softmax_type):
+        attn_weights = scale_factor * x
+        if softmax_type != SoftmaxType.SCALED:
+            big_neg = jnp.finfo(x.dtype).min
+            attn_weights = jnp.where(mask, attn_weights, big_neg)
+        result = jax.nn.softmax(attn_weights).astype(x.dtype)
+        return jnp.mean(result)
+
+    @staticmethod
+    def custom_self_fused_attn(qkv, bias, mask, rng_key, dropout_prob,
+                               attn_bias_type, attn_mask_type,
+                               scaling_factor, sharding_type):
+        mask = (mask == 0)  # invert mask
+        bias_ = None if attn_bias_type == AttnBiasType.NO_BIAS else bias
+        result = self_fused_attn(qkv, bias_, mask,
+                                 seed=rng_key,
+                                 attn_bias_type=attn_bias_type,
+                                 attn_mask_type=attn_mask_type,
+                                 scaling_factor=scaling_factor,
+                                 dropout_probability=dropout_prob,
+                                 is_training=True,
+                                 sharding_type=sharding_type)
+        return jnp.mean(result)
+
+    def reference_self_fused_attn(self, qkv, bias, mask, rng_key, dropout_prob,
+                                  attn_bias_type, attn_mask_type,
+                                  scaling_factor):
+        # split interleaved QKV into separate matrices
+        query, key, value = jnp.split(qkv, [1, 2], axis=-3)
+        return self.fused_attn_core(
+            query, key, value, bias, mask, scaling_factor,
+            attn_bias_type, attn_mask_type,
+            rng_key, dropout_prob)
+
+    @staticmethod
+    def custom_cross_fused_attn(query, key_value, mask, rng_key, dropout_prob,
+                                attn_mask_type, scaling_factor, sharding_type):
+        mask = (mask == 0)  # invert mask
+        result = cross_fused_attn(query, key_value, mask,
+                                  seed=rng_key,
+                                  attn_bias_type=AttnBiasType.NO_BIAS,
+                                  attn_mask_type=attn_mask_type,
+                                  scaling_factor=scaling_factor,
+                                  dropout_probability=dropout_prob,
+                                  is_training=True,
+                                  sharding_type=sharding_type)
+        return jnp.mean(result)
+
+    def reference_cross_fused_attn(self, query, key_value, mask, rng_key, dropout_prob,
+                                   attn_mask_type, scaling_factor):
+        key, value = jnp.split(key_value, [1], axis=-3)
+        return self.fused_attn_core(
+            query, key, value, None, mask, scaling_factor,
+            AttnBiasType.NO_BIAS, attn_mask_type,
+            rng_key, dropout_prob)

tests/jax/sharding_configs.py

+# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+import jax
+from itertools import product
+from transformer_engine.jax.sharding import ShardingType
+from transformer_engine.jax.softmax import SoftmaxType
+from transformer_engine.jax.fused_attn import AttnBiasType, AttnMaskType\
+
+
+class ShardingConfigs(object):
+    
+    def __init__(self, num_gpus=len(jax.devices())):
+        super().__init__()
+        if num_gpus < 2:
+            raise ValueError(f"ShardingConfig: Need at least 2 GPUs.")
+        
+        self.device_count = min(num_gpus, 8)
+        mesh_configs = [
+            ((self.device_count, 1),    ("dp", None), ShardingType.DP),
+            ((self.device_count, 1),    ("tp", None), ShardingType.TP_COL),
+            ((self.device_count, 1),    ("tp", None), ShardingType.TP_ROW)
+        ]
+        if self.device_count >= 4:
+            mesh_configs += [
+                ((self.device_count//2, 2), ("dp", "tp"), ShardingType.DP_TP_COL),
+                ((self.device_count//2, 2), ("dp", "tp"), ShardingType.DP_TP_ROW),
+            ]
+        if self.device_count >= 6:
+            mesh_configs += [
+                ((2, self.device_count//2), ("dp", "tp"), ShardingType.DP_TP_COL),
+                ((2, self.device_count//2), ("dp", "tp"), ShardingType.DP_TP_ROW),
+            ]
+
+        layernorm_collectives = {
+            ShardingType.DP        : {'all-reduce': 2, 'other': 0},
+            ShardingType.TP_COL    : {'all-reduce': 0, 'other': 0},
+            ShardingType.DP_TP_COL : {'all-reduce': 2, 'other': 0}
+        }
+        self.layernorm_refs = [
+            mesh_config + (layernorm_collectives[mesh_config[2]], ) \
+                for mesh_config in mesh_configs \
+                    if mesh_config[2] not in (ShardingType.TP_ROW, ShardingType.DP_TP_ROW)
+        ]
+
+        self.softmax_types = [
+            SoftmaxType.SCALED,
+            SoftmaxType.SCALED_MASKED,
+            SoftmaxType.SCALED_UPPER_TRIANG_MASKED
+        ]
+        softmax_collectives = {
+            ShardingType.DP        : {'all-reduce': 1, 'other': 0},
+            ShardingType.TP_COL    : {'all-reduce': 1, 'other': 0},
+            ShardingType.TP_ROW    : {'all-reduce': 1, 'other': 0},
+            ShardingType.DP_TP_COL : {'all-reduce': 1, 'other': 0},
+            ShardingType.DP_TP_ROW : {'all-reduce': 1, 'other': 0}
+        }
+        self.softmax_refs = [
+            mesh_config + (softmax_collectives[mesh_config[2]], ) for mesh_config in mesh_configs
+        ]
+
+        self.self_attn_bias_types = [
+            AttnBiasType.NO_BIAS,
+            AttnBiasType.PRE_SCALE_BIAS,
+            AttnBiasType.POST_SCALE_BIAS
+        ]
+        self.self_attn_mask_types = [
+            AttnMaskType.CAUSAL_MASK,
+            AttnMaskType.PADDING_MASK,
+            AttnMaskType.NO_MASK
+        ]
+        self_attn_collectives = {
+            ShardingType.DP : {
+                AttnBiasType.NO_BIAS         : {'all-reduce': 1, 'other': 0},
+                AttnBiasType.PRE_SCALE_BIAS  : {'all-reduce': 2, 'other': 0},
+                AttnBiasType.POST_SCALE_BIAS : {'all-reduce': 2, 'other': 0},
+            },
+            ShardingType.TP_COL : {
+                AttnBiasType.NO_BIAS         : {'all-reduce': 1, 'other': 0},
+                AttnBiasType.PRE_SCALE_BIAS  : {'all-reduce': 1, 'other': 0},
+                AttnBiasType.POST_SCALE_BIAS : {'all-reduce': 1, 'other': 0}
+            },
+            ShardingType.DP_TP_COL : {
+                AttnBiasType.NO_BIAS         : {'all-reduce': 1, 'other': 0},
+                AttnBiasType.PRE_SCALE_BIAS  : {'all-reduce': 2, 'other': 0},
+                AttnBiasType.POST_SCALE_BIAS : {'all-reduce': 2, 'other': 0}
+            },
+        }
+        self.self_attn_refs = [
+            mesh_config + (bias_type, self_attn_collectives[mesh_config[2]][bias_type]) \
+                for mesh_config, bias_type in product(mesh_configs, self.self_attn_bias_types) \
+                    if mesh_config[2] not in (ShardingType.TP_ROW, ShardingType.DP_TP_ROW)
+        ]
+
+        self.cross_attn_mask_types = [
+            AttnMaskType.PADDING_MASK,
+            AttnMaskType.NO_MASK
+        ]
+        self.cross_attn_refs = [
+            mesh_config + ({'all-reduce': 1, 'other': 0}, ) \
+                for mesh_config in mesh_configs \
+                    if mesh_config[2] not in (ShardingType.TP_ROW, ShardingType.DP_TP_ROW)
+        ]

tests/jax/test_distributed_custom_ops.py

+# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+import pytest
+import numpy as np
+from functools import partial
+
+import jax
+import jax.numpy as jnp
+from jax import random
+from jax.sharding import NamedSharding
+
+from utils import is_devices_enough
+from sharding_configs import *
+from custom_ops_helper import *
+from transformer_engine.jax.sharding import global_shard_guard
+from transformer_engine.jax.fused_attn import AttnBiasType, AttnMaskType
+
+configs = ShardingConfigs()   # default device count is len(jax.devices())
+helper = CustomOpsTestHelper()
+
+@pytest.mark.skipif(not helper.use_custom_partitioning(),
+                    reason='TE/JAX version does not support sharding with ' + \
+                           'jax.experimental.custom_partitioning.')
+@pytest.mark.skipif(not is_devices_enough(configs.device_count), reason='Num of GPU is not enough')
+class TestCustomPartitioningOpsGenerator:
+
+    @pytest.mark.parametrize('mesh_shape, mesh_names, sharding_type, collective_ref',
+                             configs.layernorm_refs)
+    @pytest.mark.parametrize('zero_centered_gamma', [False, True])
+    def test_layernorm(self, mesh_shape, mesh_names, sharding_type, collective_ref,
+                       zero_centered_gamma):
+        epsilon = 1e-6
+
+        custom_func = partial(helper.custom_layernorm,
+                              zero_centered_gamma=zero_centered_gamma,
+                              epsilon=epsilon,
+                              sharding_type=sharding_type)
+        
+        reference_func = partial(helper.reference_layernorm,
+                                 zero_centered_gamma=zero_centered_gamma,
+                                 epsilon=epsilon)
+
+        batch_size, _, num_heads, head_dim = helper.qkv_shape
+        hidden_size = num_heads*head_dim
+        input_shape = (batch_size, hidden_size)
+        other_shape = (hidden_size, )
+        x_ = random.normal(random.PRNGKey(1124), input_shape, dtype=helper.dtype)
+        gamma_ = jnp.ones(other_shape, dtype=helper.dtype)
+        beta_ = jnp.ones(other_shape, dtype=helper.dtype)
+
+        x_spec, gamma_spec, beta_spec = helper.get_sharding_spec(mesh_names, sharding_type)
+        devices = np.asarray(jax.devices()[:configs.device_count]).reshape(*mesh_shape)
+        mesh = jax.sharding.Mesh(devices, mesh_names)
+        with mesh, global_shard_guard(helper.get_sharding_resource(mesh_names, sharding_type)):
+            x_ = jax.device_put(x_, NamedSharding(mesh, x_spec))
+            gamma_ = jax.device_put(gamma_, NamedSharding(mesh, gamma_spec))
+            beta_ = jax.device_put(beta_, NamedSharding(mesh, beta_spec))
+            helper.compare_ops(
+                custom_func, reference_func, collective_ref,
+                x_, gamma_, beta_, grad_args=(0, 1, 2), dtype=helper.dtype,
+                in_shardings=[x_spec, gamma_spec, beta_spec],
+                out_shardings=(None, (x_spec, gamma_spec, beta_spec))
+            )
+
+    @pytest.mark.parametrize('mesh_shape, mesh_names, sharding_type, collective_ref',
+                             configs.layernorm_refs)
+    def test_rmsnorm(self, mesh_shape, mesh_names, sharding_type, collective_ref):
+        epsilon = 1e-6
+        custom_func = partial(helper.custom_rmsnorm, epsilon=epsilon,sharding_type=sharding_type)
+        reference_func = partial(helper.reference_rmsnorm, epsilon=epsilon)
+
+        batch_size, _, num_heads, head_dim = helper.qkv_shape
+        hidden_size = num_heads*head_dim
+        input_shape = (batch_size, hidden_size)
+        other_shape = (hidden_size, )
+        x_ = random.normal(random.PRNGKey(1124), input_shape, dtype=helper.dtype)
+        gamma_ = jnp.ones(other_shape, dtype=helper.dtype)
+
+        x_spec, gamma_spec = helper.get_sharding_spec(mesh_names, sharding_type)
+        devices = np.asarray(jax.devices()[:configs.device_count]).reshape(*mesh_shape)
+        mesh = jax.sharding.Mesh(devices, mesh_names)
+        with mesh, global_shard_guard(helper.get_sharding_resource(mesh_names, sharding_type)):
+            x_ = jax.device_put(x_, NamedSharding(mesh, x_spec))
+            gamma_ = jax.device_put(gamma_, NamedSharding(mesh, gamma_spec))
+            helper.compare_ops(
+                custom_func, reference_func, collective_ref,
+                x_, gamma_, grad_args=(0, 1), dtype=helper.dtype,
+                in_shardings=[x_spec, gamma_spec],
+                out_shardings=(None, (x_spec, gamma_spec))
+            )
+
+    @pytest.mark.parametrize('mesh_shape, mesh_names, sharding_type, collective_ref',
+                             configs.softmax_refs)
+    @pytest.mark.parametrize('softmax_type', configs.softmax_types)
+    def test_softmax(self, mesh_shape, mesh_names, sharding_type, collective_ref,
+                     softmax_type):
+        batch_size, seq_len, num_heads, head_dim = helper.qkv_shape
+        scale_factor = 1./jnp.sqrt(head_dim)
+        
+        custom_func = partial(helper.custom_softmax,
+                              scale_factor=scale_factor,
+                              softmax_type=softmax_type,
+                              sharding_type=sharding_type)
+        reference_func = partial(helper.reference_softmax,
+                                 scale_factor=scale_factor,
+                                 softmax_type=softmax_type)
+
+        input_size = (batch_size, num_heads, seq_len, seq_len)
+        x_ = random.normal(random.PRNGKey(1124), input_size, dtype=helper.dtype)
+
+        pad_len = int(seq_len * helper.pad_ratio)
+        valid_len = seq_len - pad_len
+        tokens = jnp.concatenate((jnp.ones((batch_size, valid_len), dtype=helper.dtype),
+                                  jnp.zeros((batch_size, pad_len), dtype=helper.dtype)),
+                                 axis=-1)
+        mask_ = helper.make_mask(tokens, tokens, AttnMaskType.PADDING_MASK)
+
+        x_spec, mask_spec = helper.get_sharding_spec(mesh_names, sharding_type)
+        devices = np.asarray(jax.devices()[:configs.device_count]).reshape(*mesh_shape)
+        mesh = jax.sharding.Mesh(devices, mesh_names)
+        with mesh, global_shard_guard(helper.get_sharding_resource(mesh_names, sharding_type)):
+            x_ = jax.device_put(x_, NamedSharding(mesh, x_spec))
+            mask_ = jax.device_put(mask_, NamedSharding(mesh, mask_spec))
+            helper.compare_ops(
+                custom_func, reference_func, collective_ref,
+                (0), x_, mask_, grad_args=(0), dtype=helper.dtype,
+                in_shardings=[x_spec, mask_spec],
+                out_shardings=(None, (x_spec))
+            )
+
+    @pytest.mark.parametrize(
+        'mesh_shape, mesh_names, sharding_type, attn_bias_type, collective_ref',
+        configs.self_attn_refs)
+    @pytest.mark.parametrize('attn_mask_type', configs.self_attn_mask_types)
+    def test_self_fused_attn(self, mesh_shape, mesh_names, sharding_type, collective_ref,
+                             attn_bias_type, attn_mask_type, backend):
+        batch_size, seq_len, num_heads, head_dim = helper.qkv_shape
+        helper.check_fused_attn_inputs(seq_len, seq_len, head_dim,
+                                       helper.pad_ratio, helper.dropout_prob,
+                                       attn_bias_type, attn_mask_type, backend)
+
+        dropout_rng = random.PRNGKey(91023051)
+        split_rng = random.split(dropout_rng, configs.device_count)
+        scale_factor = 1./jnp.sqrt(head_dim)
+
+        custom_func = partial(helper.custom_self_fused_attn,
+                              rng_key=split_rng,
+                              dropout_prob=helper.dropout_prob,
+                              attn_bias_type=attn_bias_type,
+                              attn_mask_type=attn_mask_type,
+                              scaling_factor=scale_factor,
+                              sharding_type=sharding_type)
+        reference_func = partial(helper.reference_self_fused_attn,
+                                 rng_key=dropout_rng,
+                                 dropout_prob=helper.dropout_prob,
+                                 attn_bias_type=attn_bias_type,
+                                 attn_mask_type=attn_mask_type,
+                                 scaling_factor=scale_factor)
+
+        key = random.PRNGKey(1124)
+        subkeys = random.split(key, 2)
+
+        qkv_shape = (batch_size, seq_len, 3, num_heads, head_dim)
+        qkv_ = random.normal(subkeys[0], qkv_shape, dtype=helper.dtype)
+        bias_shape = (1, num_heads, seq_len, seq_len)
+        bias_ = random.normal(subkeys[1], bias_shape, dtype=helper.dtype)
+        
+        pad_len = int(seq_len * helper.pad_ratio)
+        valid_len = seq_len - pad_len
+        tokens = jnp.concatenate((jnp.ones((batch_size, valid_len), dtype=helper.dtype),
+                                  jnp.zeros((batch_size, pad_len), dtype=helper.dtype)),
+                                 axis=-1)
+        mask_ = helper.make_mask(tokens, tokens, attn_mask_type)
+
+        qkv_spec, bias_spec, mask_spec = helper.get_sharding_spec(mesh_names, sharding_type)
+        devices = np.asarray(jax.devices()[:configs.device_count]).reshape(*mesh_shape)
+        mesh = jax.sharding.Mesh(devices, mesh_names)
+        with mesh, global_shard_guard(helper.get_sharding_resource(mesh_names, sharding_type)):
+            qkv_ = jax.device_put(qkv_, NamedSharding(mesh, qkv_spec))
+            bias_ = jax.device_put(bias_, NamedSharding(mesh, bias_spec))
+            mask_ = jax.device_put(mask_, NamedSharding(mesh, mask_spec))
+            helper.compare_ops(
+                custom_func, reference_func, collective_ref,
+                qkv_, bias_, mask_, grad_args=(0, 1), dtype=helper.dtype,
+                in_shardings=[qkv_spec, bias_spec, mask_spec],
+                out_shardings=(None, (qkv_spec, bias_spec))
+            )
+
+    @pytest.mark.parametrize('mesh_shape, mesh_names, sharding_type, collective_ref',
+                             configs.cross_attn_refs)
+    @pytest.mark.parametrize('attn_mask_type', configs.cross_attn_mask_types)
+    def test_cross_fused_attn(self, mesh_shape, mesh_names, sharding_type, collective_ref,
+                              attn_mask_type, backend):
+        batch_size, seq_len, num_heads, head_dim = helper.qkv_shape
+        helper.check_fused_attn_inputs(seq_len, seq_len, head_dim,
+                                       helper.pad_ratio, helper.dropout_prob,
+                                       AttnBiasType.NO_BIAS, attn_mask_type, backend)
+        
+        dropout_rng = random.PRNGKey(91023051)
+        split_rng = random.split(dropout_rng, configs.device_count)
+        scale_factor = 1./jnp.sqrt(head_dim)
+
+        custom_func = partial(helper.custom_cross_fused_attn,
+                              rng_key=split_rng,
+                              dropout_prob=helper.dropout_prob,
+                              attn_mask_type=attn_mask_type,
+                              scaling_factor=scale_factor,
+                              sharding_type=sharding_type)
+        reference_func = partial(helper.reference_cross_fused_attn,
+                                 rng_key=split_rng,
+                                 dropout_prob=helper.dropout_prob,
+                                 attn_mask_type=attn_mask_type,
+                                 scaling_factor=scale_factor)
+
+        key = random.PRNGKey(1124)
+        subkeys = random.split(key, 2)
+
+        q_shape = (batch_size, seq_len, num_heads, head_dim)
+        q_ = random.normal(subkeys[0], q_shape, dtype=helper.dtype)
+        kv_shape = (batch_size, seq_len, 2, num_heads, head_dim)
+        kv_ = random.normal(subkeys[1], kv_shape, dtype=helper.dtype)
+        
+        pad_len = int(seq_len * helper.pad_ratio)
+        valid_len = seq_len - pad_len
+        tokens = jnp.concatenate((jnp.ones((batch_size, valid_len), dtype=helper.dtype),
+                                  jnp.zeros((batch_size, pad_len), dtype=helper.dtype)),
+                                 axis=-1)
+        mask_ = helper.make_mask(tokens, tokens, attn_mask_type)
+
+        q_spec, kv_spec, mask_spec = helper.get_sharding_spec(mesh_names, sharding_type)
+        devices = np.asarray(jax.devices()[:configs.device_count]).reshape(*mesh_shape)
+        mesh = jax.sharding.Mesh(devices, mesh_names)
+        with mesh, global_shard_guard(helper.get_sharding_resource(mesh_names, sharding_type)):
+            q_ = jax.device_put(q_, NamedSharding(mesh, q_spec))
+            kv_= jax.device_put(kv_, NamedSharding(mesh, kv_spec))
+            mask_ = jax.device_put(mask_, NamedSharding(mesh, mask_spec))
+            helper.compare_ops(
+                custom_func, reference_func, collective_ref,
+                q_, kv_, mask_, grad_args=(0, 1), dtype=helper.dtype,
+                in_shardings=[q_spec, kv_spec, mask_spec],
+                out_shardings=(None, (q_spec, kv_spec))
+            )