# Copyright 2023 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
Note: Mixture of Expert (MoE) structure. This is an experimental interface that is subject to change or deletion.
"""
from __future__ import absolute_import
from __future__ import division

import math
import copy
import numpy as np
import mindspore.ops as ops

from mindspore.common.tensor import Tensor
from mindspore.common.parameter import Parameter
from mindspore.common.initializer import initializer, Normal
from mindspore.ops.operations._sequence_ops import TensorToScalar
import mindspore.common.dtype as mstype

# MindSpore 2.0 has changed the APIs of _checkparam, the following try except is for compatibility
try:
    from mindspore._checkparam import Validator
except ImportError:
    import mindspore._checkparam as Validator
from mindspore.ops import operations as P
from mindspore.ops import functional as F
from mindspore.ops.primitive import constexpr
from mindspore.ops.auto_generate import MoeFinalizeRouting
from mindspore.nn.cell import Cell
from mindspore.nn.layer import Dense
from mindspore.context import ParallelMode
from mindspore.parallel._utils import _get_parallel_mode, _is_sharding_propagation

from mindformers.modules.transformer.op_parallel_config import default_moeparallel_config
from mindformers.modules.transformer.moe_utils import ZLoss
from mindformers.tools.utils import get_predict_run_mode

__all__ = [
    "MoEConfig"]

dtype_map = {
    'float16': mstype.float32,
    'float32': mstype.float32,
    'bfloat16': mstype.bfloat16
}


def _check_aux_loss_config(aux_loss_types, aux_loss_factors):
    """
        Check if aux_loss_types and aux_loss_factors are valid.
        Args:
            aux_loss_types: list of auxiliary loss types
            aux_loss_factors: list of auxiliary loss factors

        Returns:
            aux_loss_config (dict): dict of auxiliary loss types and factors.
    """
    supported_loss_types = ["expert", "device", "comm"]

    if aux_loss_types is None:
        aux_loss_types = []
        aux_loss_factors = []
    else:
        if not (isinstance(aux_loss_types, list) and isinstance(aux_loss_factors, list)):
            raise ValueError(f"Auxiliary loss types and factors should be list, bug got {aux_loss_types} and "
                             f"{aux_loss_factors}")
    if set(aux_loss_types) - set(supported_loss_types):
        raise ValueError(f"Auxiliary loss types in {supported_loss_types} only supported, but got {aux_loss_types}")
    if aux_loss_factors is None:
        raise ValueError(f"Got auxiliary loss types {aux_loss_types}, but corresponding loss factors are not set.")

    return aux_loss_types, aux_loss_factors


class MoEConfig:
    r"""
        The configuration of MoE (Mixture of Expert).

        Args:
            expert_num (int): The number of experts employed. Default: 1
            capacity_factor (float): The factor is used to indicate how much to expand expert capacity,
                which is >=1.0. Default: 1.1.
            aux_loss_factor (float): The factor is used to indicate how much the load balance loss (produced by the
                router) to be added to the entire model loss, which is < 1.0. Default: 0.05.
            num_experts_chosen (int): The number of experts is chosen by each token and it should not be larger
                than expert_num. Default: 1.
            expert_group_size (int): The number of tokens in each data parallel group. Default: None. This parameter is
                effective only when in AUTO_PARALLEL mode, and NOT SHARDING_PROPAGATION.
            group_wise_a2a (bool): Whether to enable group-wise alltoall communication, which can reduce communication
                time by converting part of intercommunication into intra communication. Default: False. This parameter
                is effective only when model parallel > 1 and data_parallel equal to expert parallel.
            comp_comm_parallel (bool): Whether to enable ffn compute and communication parallel, which can reduce pure
                communicattion time by splitting and overlapping compute and communication. Default: False.
            comp_comm_parallel_degree (int): The split number of compute and communication. The larger the numbers,
                the more overlap there will be but will consume more memory. Default: 2. This parameter is effective
                only when comp_comm_parallel enable.
            routing_policy (str): The routing policy to use in MoE layer. Default: TopkRouterV1.
            moe_shared_expert_overlap (bool): Whether to enable shared expert parallel, which can overlap
                the computing of the shared expert and the communication of the routing expert. Default: False.
            use_gmm (bool): Whether to enable MOVE3 module, which uses gmm op instead of bmm. Default: False.
            enable_gmm_safe_tokens (bool): Whether to pad safe token. If use MOVE3 module,
                when an expert accept 0 token, gmm will report error. Default: False.
            dispatch_global_max_bs (int): In decode of infer, it should be 0 or global max batch_size * seq_len(1)
                for dispatch ops and combine ops. Default: 0.
            ep_extend_tp (bool): Whether to enhances the parallel computing efficiency by combining
                Expert Parallelism (EP) and Tensor Parallelism (TP). Default: True.

        Supported Platforms:
            ``Ascend`` ``GPU``

        Examples:
            >>> from mindformers.modules.transformer import MoEConfig
            >>> moe_config = MoEConfig(expert_num=4, capacity_factor=5.0, aux_loss_factor=0.05, num_experts_chosen=1,
            ...                        expert_group_size=64, group_wise_a2a=True, comp_comm_parallel=False,
            ...                        comp_comm_parallel_degree=2, routing_policy="TopkRouterV2")
    """

    def __init__(self, expert_num=1, capacity_factor=1.1, aux_loss_factor=0.05, num_experts_chosen=1,
                 expert_group_size=None, group_wise_a2a=False, comp_comm_parallel=False, comp_comm_parallel_degree=2,
                 save_token_distribution=False, cur_layer=0, enable_cold_hot_expert=False, update_step=10000,
                 hot_expert_num=0, cold_token_percent=1.0, moe_module_name="", routing_policy="TopkRouterV1",
                 norm_topk_prob=True, enable_sdrop=False, use_fused_ops_topkrouter=False, router_dense_type="float32",
                 shared_expert_num=0, use_shared_expert_gating=False, max_router_load=128 * 1024,
                 topk_method="greedy", topk_group=None, n_group=None, enable_deredundency=False, npu_nums_per_device=1,
                 first_k_dense_replace=True, moe_intermediate_size=1407, routed_scaling_factor=1.0,
                 aux_loss_types=None, aux_loss_factors=None, z_loss_factor=0., balance_via_topk_bias=False,
                 topk_bias_update_rate=0., use_allgather_dispatcher=False, moe_shared_expert_overlap=False,
                 expert_model_parallel=None, use_gating_sigmoid=False, use_gmm=False, enable_gmm_safe_tokens=False,
                 use_fused_ops_permute=False, callback_moe_droprate=False, dispatch_global_max_bs=0, ep_extend_tp=True):
        Validator.check_positive_int(expert_num, "expert_num")
        Validator.check_positive_float(aux_loss_factor, "aux_loss_factor")
        Validator.check_positive_int(num_experts_chosen, "num_experts_chosen")
        Validator.check_bool(group_wise_a2a, "group_wise_a2a")
        Validator.check_bool(comp_comm_parallel, "comp_comm_parallel")
        Validator.check_positive_int(comp_comm_parallel_degree, "comp_comm_parallel_degree")
        Validator.check_bool(save_token_distribution, "save_token_distribution")
        Validator.check_non_negative_int(cur_layer, "cur_layer")
        Validator.check_bool(enable_cold_hot_expert, "enable_cold_hot_expert")
        Validator.check_positive_int(update_step, "update_step")
        Validator.check_non_negative_int(hot_expert_num, "hot_expert_num")
        Validator.check_non_negative_float(cold_token_percent, "cold_token_percent")
        Validator.check_string(router_dense_type, ["float16", "float32", "bfloat16"], "router_dense_type")
        Validator.check_bool(ep_extend_tp, "ep_extend_tp")
        if expert_group_size is not None:
            Validator.check_positive_int(expert_group_size, "expert_group_size")
        if aux_loss_factor >= 1.0:
            raise ValueError(f"'aux_loss_factor' must be less than 1.0, "
                             f"but got {aux_loss_factor}.")
        if num_experts_chosen > expert_num:
            raise ValueError(f"'num_experts_chosen' must not be larger than 'expert_num', "
                             f"but got {num_experts_chosen}.")
        if hot_expert_num > expert_num:
            raise ValueError(f"'hot_expert_num' must not be larger than 'expert_num', "
                             f"but got {hot_expert_num}.")
        if cold_token_percent > 1.0 or cold_token_percent <= 0.0:
            raise ValueError(f"'cold_token_percent' must be in the range (0.0, 1.0], "
                             f"but got {cold_token_percent}.")

        self.expert_num = expert_num
        self.capacity_factor = capacity_factor
        self.aux_loss_factor = aux_loss_factor
        self.num_experts_chosen = num_experts_chosen
        self.expert_group_size = expert_group_size
        self.group_wise_a2a = group_wise_a2a
        self.comp_comm_parallel = comp_comm_parallel
        self.comp_comm_parallel_degree = comp_comm_parallel_degree
        self.save_token_distribution = save_token_distribution
        self.cur_layer = cur_layer
        self.enable_cold_hot_expert = enable_cold_hot_expert
        self.update_step = update_step
        self.hot_expert_num = hot_expert_num
        self.cold_token_percent = cold_token_percent
        self.moe_module_name = moe_module_name
        self.routing_policy = routing_policy
        self.norm_topk_prob = norm_topk_prob
        self.enable_sdrop = enable_sdrop
        self.use_fused_ops_topkrouter = use_fused_ops_topkrouter
        self.router_dense_type = dtype_map.get(router_dense_type)
        self.shared_expert_num = shared_expert_num
        self.use_shared_expert_gating = use_shared_expert_gating
        self.routed_scaling_factor = routed_scaling_factor
        self.moe_intermediate_size = moe_intermediate_size
        self.n_group = n_group
        self.topk_group = topk_group
        self.topk_method = topk_method
        self.first_k_dense_replace = first_k_dense_replace
        self.aux_loss_types, self.aux_loss_factors = _check_aux_loss_config(aux_loss_types, aux_loss_factors)
        self.z_loss_factor = z_loss_factor
        self.max_router_load = max_router_load
        self.balance_via_topk_bias = balance_via_topk_bias
        self.topk_bias_update_rate = topk_bias_update_rate
        self.use_allgather_dispatcher = use_allgather_dispatcher
        self.moe_shared_expert_overlap = moe_shared_expert_overlap
        self.expert_model_parallel = expert_model_parallel
        self.use_gating_sigmoid = use_gating_sigmoid
        self.enable_deredundency = enable_deredundency
        self.npu_nums_per_device = npu_nums_per_device
        self.use_gmm = use_gmm
        self.enable_gmm_safe_tokens = enable_gmm_safe_tokens
        self.use_fused_ops_permute = use_fused_ops_permute
        self.callback_moe_droprate = callback_moe_droprate
        self.dispatch_global_max_bs = dispatch_global_max_bs

    def __eq__(self, other) -> bool:
        return isinstance(other, MoEConfig) and (self.to_dict() == other.to_dict())

    def to_diff_dict(self):
        """
        Compare the configuration dictionary of the current object with the default configuration dictionary,
        identify the differences between the two, and store these differences in a new dictionary called res-dict
        """
        config_dict = self.to_dict()
        default_dict = MoEConfig().to_dict()
        res_dict = {}
        for k, v in config_dict.items():
            if v != default_dict.get(k):
                res_dict[k] = v
        return res_dict

    def to_dict(self):
        return copy.deepcopy(self.__dict__)


default_moe_config = MoEConfig()


@constexpr
def calculate_expert_capacity(k, tokens_per_group, capacity_factor, expert_dim):
    return math.ceil(k * tokens_per_group * capacity_factor / expert_dim)


@constexpr
def calculate_expert_capacity_v2(k, tokens_per_group, capacity_factor, expert_dim, mp):
    raw_capacity = math.ceil(k * tokens_per_group * capacity_factor / expert_dim)
    if tokens_per_group < raw_capacity:
        raw_capacity = tokens_per_group
    if raw_capacity % mp > 0:
        raw_capacity = raw_capacity + mp - (raw_capacity % mp)
    return raw_capacity


class MoEV2(Cell):
    """
    The mixture of experts (MoE) implementation. The implementation includes a router and a FeedForward layer.
    The router dispatches tokens to experts in FeedForward, then FeedForward does computation, and the final output is
    obtained by multiplying FeedForward's output and router's combine weight.
    This is a common interface, which allows any ffn class and any Router algorithm(implemented in V2 form).

    Args:
        hidden_size (int): The dimension of the inputs.
        ffn_hidden_size (int): The intermediate hidden size.
        dropout_rate (float): The dropout rate for the second linear's output.
        hidden_act (str): The activation of the internal feedforward layer. Supports 'relu',
                         'relu6', 'tanh', 'gelu', 'fast_gelu', 'elu', 'sigmoid', 'prelu', 'leakyrelu', 'hswish',
                         'hsigmoid', 'logsigmoid' and so on. Default: gelu.
        param_init_type (dtype.Number): The parameter initialization type. Can be dtype.float32 or dtype.float16.
        moe_config(MoEConfig): The configuration of MoE (Mixture of Expert). Default is an instance of MoEConfig with
            default values. Please see `MoEConfig`.
        parallel_config(MoEParallelConfig): The parallel config for MoE, see `MoEParallelConfig`.
            Default `default_moeparallel_config`, an instance of `MoEParallelConfig` with default args.
        return_extra_loss (bool): whether to return extra regularization loss. Default False.

    Inputs:
        - **x** (Tensor) - should be `[batch, seq_length, hidden_size]`. Float tensor.

    Outputs:
        Tensor, the output of this layer after mapping. The shape is `[batch, seq_length, hidden_size]`.
    """

    def __init__(self,
                 ffn,
                 dim,
                 moe_config=default_moe_config,
                 parallel_config=default_moeparallel_config,
                 return_extra_loss=False,
                 init_method_std=0.01):
        super(MoEV2, self).__init__()
        self.hidden_size = dim
        self.expert_dim = moe_config.expert_num
        self.return_extra_loss = return_extra_loss
        self.capacity_factor = moe_config.capacity_factor
        self.num_experts_chosen = moe_config.num_experts_chosen
        self.dp_group = parallel_config.data_parallel * parallel_config.context_parallel
        self.dp = parallel_config.data_parallel * parallel_config.context_parallel
        self.ep = parallel_config.expert_parallel
        self.mp = parallel_config.model_parallel
        self.use_allgather_dispatcher = moe_config.use_allgather_dispatcher
        self.group_wise_a2a = moe_config.group_wise_a2a if self.mp > 1 else False
        self.add_loss = P.Add()

        self.ffn = ffn
        Validator.check_string(moe_config.routing_policy, ["TopkRouterV2"], "routing_policy")
        self.ffn_forward = self._ffn_forward
        router_parallel_config = copy.deepcopy(parallel_config)
        self.mp_moe = moe_config.expert_model_parallel if moe_config.expert_model_parallel is not None \
            else parallel_config.model_parallel
        if moe_config.expert_model_parallel is not None:
            self.dp = self.dp * self.mp // self.mp_moe
            self.dp_group = self.dp
            self.mp = self.mp_moe
            router_parallel_config.data_parallel = self.dp
            router_parallel_config.model_parallel = self.mp
            router_parallel_config.context_parallel = 1
        self.dp_moe = self.dp // self.ep
        self.dp_range = Tensor(np.arange(self.dp_group).reshape(-1, 1), mstype.int32)  # (dp, 1) = [[0],[1],[2]...[dp]]
        self.use_seq_parallel = parallel_config.use_seq_parallel
        if self.use_seq_parallel:
            self.ffn_forward = self._ffn_forward_sq
            router_parallel_config.data_parallel = self.dp * self.mp
            router_parallel_config.model_parallel = 1
            router_parallel_config.context_parallel = 1
            self.dp_group *= self.mp
        self.router = Router(d_model=self.hidden_size,
                             moe_config=moe_config,
                             routing_policy=moe_config.routing_policy,
                             training=(not get_predict_run_mode()),
                             parallel_config=router_parallel_config,
                             init_method_std=init_method_std)

        self.reshape = P.Reshape()
        self.shape = P.Shape()
        self.cast = P.Cast()
        self.transpose_4dim_dp1 = P.Transpose().shard(((1, self.dp, 1, 1),))
        self.transpose_4dim_dp0 = P.Transpose().shard(((self.dp, 1, 1, 1),))
        self.transpose_5dim_ep2 = P.Transpose().shard(((self.dp_moe, 1, self.ep, 1, 1),))
        self.concat_dp = P.Concat(2).shard(((1, self.dp, 1, 1), (1, self.dp, 1, 1)))
        self.stride_slice = P.StridedSlice().shard(((self.dp, 1, 1, 1),))
        self.stride_slice_dp = P.StridedSlice().shard(((1, self.dp, 1, 1),))
        self.stride_slice_ep = P.StridedSlice().shard(((self.ep, 1, 1, 1),))
        self.stride_slice_dp_mp = P.StridedSlice().shard(((1, self.dp, self.mp, 1),))
        self.stride_slice_ep_mp = P.StridedSlice().shard(((self.ep, 1, self.mp, 1),))

        # outer_dp groupwiseall2all
        self.stride_slice_outer_dp_mp = P.StridedSlice().shard(((self.dp_moe, 1, self.ep, self.mp, 1),))
        self.stride_slice_outer_ep_mp = P.StridedSlice().shard(((self.dp_moe, self.ep, 1, self.mp, 1),))
        self.stride_slice_outer_ep = P.StridedSlice().shard(((self.dp_moe, self.ep, 1, 1, 1),))
        self.stride_slice_outer_dp = P.StridedSlice().shard(((self.dp_moe, 1, self.ep, 1, 1),))
        self.transpose_5dim_ep1 = P.Transpose().shard(((self.dp_moe, self.ep, 1, 1, 1),))
        # outer dp
        self.transpose_mp1 = P.Transpose().shard(((self.dp_moe, self.ep, self.mp, 1, 1, 1),))
        self.transpose_mp2 = P.Transpose().shard(((self.dp_moe, 1, self.ep, self.mp, 1, 1),))
        self.stride_slice_allgather = P.StridedSlice().shard(((self.dp_moe, self.ep, 1, 1, 1),))
        # interleave
        self.comp_comm_parallel = moe_config.comp_comm_parallel
        self.comp_comm_parallel_degree = moe_config.comp_comm_parallel_degree
        if self.comp_comm_parallel:
            self.split = P.Split(axis=2, output_num=self.comp_comm_parallel_degree).shard(((self.dp_group, 1, 1, 1),))
            self.concat = P.Concat(2).shard(tuple((self.dp_group, 1, 1, 1) \
                                                  for _ in range(self.comp_comm_parallel_degree)))
            self.split.add_prim_attr("enable_interleave", 1)
            self.concat.add_prim_attr("enable_interleave", 1)

    def ffn_parallel_forward(self, expert_input, capacity):
        """
        use comp_comm_overlap Computing the FFN.
        """
        if self.comp_comm_parallel:
            sub_capacity = capacity // self.comp_comm_parallel_degree
            output_list = []
            for sub_expert_input in self.split(expert_input):
                sub_expert_output = self.ffn_forward(sub_expert_input, sub_capacity)
                output_list.append(sub_expert_output)
            expert_output = self.concat(output_list)
        else:
            expert_output = self.ffn_forward(expert_input, capacity)
        return expert_output

    def _ffn_forward_sq(self, expert_input, capacity):
        """
        use sp Computing the FFN.
        """
        # expert_input shape: (dp_moe, ep, mp, E, n, h) <-- (dp, E, n, h)
        expert_input = self.reshape(expert_input, (self.dp_moe, self.ep, self.mp,
                                                   self.expert_dim, capacity, self.hidden_size))
        # dp_moe <==> outer_dp <==> dp // ep
        # expert_input shape: (dp_moe, E, ep, mp, n, h) <-- (dp_moe, ep, mp, E, n, h)
        expert_input = self.transpose_mp1(expert_input, (0, 3, 1, 2, 4, 5))
        expert_input = self.reshape(expert_input, (self.dp_moe, self.expert_dim, self.ep,
                                                   self.mp * capacity, self.hidden_size))
        expert_input = self.stride_slice_outer_dp_mp(expert_input, (0, 0, 0, 0, 0),
                                                     (self.dp_moe, self.expert_dim, self.ep,
                                                      self.mp * capacity, self.hidden_size),
                                                     (1, 1, 1, 1, 1))
        expert_input = self.stride_slice_outer_ep_mp(expert_input, (0, 0, 0, 0, 0),
                                                     (self.dp_moe, self.expert_dim, self.ep,
                                                      self.mp * capacity, self.hidden_size),
                                                     (1, 1, 1, 1, 1))
        expert_input = self.stride_slice_allgather(expert_input, (0, 0, 0, 0, 0),
                                                   (self.dp_moe, self.expert_dim, self.ep,
                                                    self.mp * capacity, self.hidden_size),
                                                   (1, 1, 1, 1, 1))

        # ffns
        expert_input = self.reshape(expert_input, (self.dp_moe, -1, self.hidden_size))
        expert_output = self.ffn(expert_input)
        expert_output = self.reshape(expert_output, (self.dp_moe, self.expert_dim, self.ep,
                                                     self.mp * capacity, self.hidden_size))

        # all2all
        expert_output = self.stride_slice_outer_ep_mp(expert_output, (0, 0, 0, 0, 0),
                                                      (self.dp_moe, self.expert_dim, self.ep,
                                                       self.mp * capacity, self.hidden_size),
                                                      (1, 1, 1, 1, 1))
        expert_output = self.stride_slice_outer_dp_mp(expert_output, (0, 0, 0, 0, 0),
                                                      (self.dp_moe, self.expert_dim, self.ep,
                                                       self.mp * capacity, self.hidden_size),
                                                      (1, 1, 1, 1, 1))
        expert_output = self.reshape(expert_output, (self.dp_moe, self.expert_dim, self.ep,
                                                     self.mp, capacity, self.hidden_size))
        expert_output = self.transpose_mp2(expert_output, (0, 2, 3, 1, 4, 5))
        expert_output = self.reshape(expert_output, (self.dp * self.mp, self.expert_dim, -1, self.hidden_size))
        return expert_output

    def _ffn_forward(self, expert_input, capacity):
        """
        Computing the FFN.
        """
        if self.group_wise_a2a:
            # expert_input shape: (dp_moe, ep, E, n, h) <-- (dp, E, n, h)
            expert_input = self.reshape(expert_input,
                                        (self.dp_moe, self.ep, self.expert_dim, -1, self.hidden_size))
            # dp_moe <==> outer_dp <==> dp // ep
            # expert_input shape: (dp_moe, E, ep, n, h) <-- (dp_moe, ep, E, n, h)
            expert_input = self.transpose_5dim_ep1(expert_input, (0, 2, 1, 3, 4))
            # capacity shard by mp
            expert_input = self.stride_slice_outer_dp_mp(expert_input, (0, 0, 0, 0, 0),
                                                         (self.dp_moe, self.expert_dim,
                                                          self.ep, capacity, self.hidden_size),
                                                         (1, 1, 1, 1, 1))
            # group-wise alltoall
            expert_input = self.stride_slice_outer_ep_mp(expert_input, (0, 0, 0, 0, 0),
                                                         (self.dp_moe, self.expert_dim,
                                                          self.ep, capacity, self.hidden_size),
                                                         (1, 1, 1, 1, 1))
            # allgather
            expert_input = self.stride_slice_outer_ep(expert_input, (0, 0, 0, 0, 0),
                                                      (self.dp_moe, self.expert_dim,
                                                       self.ep, capacity, self.hidden_size),
                                                      (1, 1, 1, 1, 1))
        else:
            # expert_input shape: (dp_moe, ep, E, n, h) <-- (dp, E, n, h)
            expert_input = self.reshape(expert_input,
                                        (self.dp_moe, self.ep, self.expert_dim, -1, self.hidden_size))
            # expert_input shape: (dp_moe, E, ep, n, h) <-- (dp_moe, ep, E, n, h)
            expert_input = self.transpose_5dim_ep2(expert_input, (0, 2, 1, 3, 4))

        expert_input = self.reshape(expert_input, (self.dp_moe, -1, self.hidden_size))
        expert_output = self.ffn(expert_input)

        if self.group_wise_a2a:
            expert_output = self.reshape(expert_output,
                                         (self.dp_moe, self.expert_dim, self.ep, -1, self.hidden_size))
            # dp_moe <==> outer_dp <==> dp // ep
            # capacity shard by mp
            expert_output = self.stride_slice_outer_ep_mp(expert_output, (0, 0, 0, 0, 0),
                                                          (self.dp_moe, self.expert_dim,
                                                           self.ep, capacity, self.hidden_size),
                                                          (1, 1, 1, 1, 1))
            # group-wise alltoall
            expert_output = self.stride_slice_outer_dp_mp(expert_output, (0, 0, 0, 0, 0),
                                                          (self.dp_moe, self.expert_dim,
                                                           self.ep, capacity, self.hidden_size),
                                                          (1, 1, 1, 1, 1))
            # allgather
            expert_output = self.stride_slice_outer_dp(expert_output, (0, 0, 0, 0, 0),
                                                       (self.dp_moe, self.expert_dim,
                                                        self.ep, capacity, self.hidden_size),
                                                       (1, 1, 1, 1, 1))
            expert_output = self.transpose_5dim_ep2(expert_output, (0, 2, 1, 3, 4))
        else:
            expert_output = self.reshape(expert_output,
                                         (self.dp_moe, self.expert_dim, self.ep, -1, self.hidden_size))
            # expert_output shape: (dp_moe, ep, E, n, h) <-- (dp_moe, E, ep, n, h)
            expert_output = self.transpose_5dim_ep2(expert_output, (0, 2, 1, 3, 4))
        expert_output = self.reshape(expert_output, (self.dp, self.expert_dim, -1, self.hidden_size))
        return expert_output

    def construct(self, input_tensor, extra_loss=0., seq_chunk=None):
        """forward process"""
        input_tensor_shape = self.shape(input_tensor)
        input_tensor = self.reshape(input_tensor, (self.dp_group, -1, self.hidden_size))  # (dp, N, h) <-- (B*S, h)

        # calculate router, we do not use router_aux_loss right now
        # (dp, E, n)int32, (dp, N, k)int32, (dp, N, k)fp16, (1,) <-- (dp, N, h),
        # where 0<= dispatch_index < 1+N, 0<= combine_index <E*(1+n)
        dispatch_policy, combine_policy, router_coeff, router_aux_loss = self.router(input_tensor, seq_chunk=seq_chunk)

        # dispatch
        expert_capacity = dispatch_policy.shape[-1]
        # expert_input shape: (dp, E, n, h) <-- (dp, N, h), (dp, E, n)
        expert_input = self.router.router.dispatch(input_tensor, dispatch_policy)
        if self.use_allgather_dispatcher:
            expert_output = self.ffn(expert_input)
        else:
            # ffn shape: (E, dp, n, h) <-- (E, dp, n, h)
            expert_output = self.ffn_parallel_forward(expert_input, expert_capacity)
        # combine shape: (dp, N, k, h) <-- (dp, E*(1+n), h), (dp, N, k)
        output_tensor = self.router.router.combine(expert_output, combine_policy, router_coeff)
        output_tensor = self.reshape(output_tensor, input_tensor_shape)  # (B*S, h) <-- (dp, N, h)
        if self.return_extra_loss:
            final_extra_loss = self.add_loss(extra_loss, router_aux_loss)
            return output_tensor, final_extra_loss
        return output_tensor  # (dp, N, h)


class Router(Cell):
    r"""
        A router backbone used to calculate logits of each token, which should be cascaded by router implementations
        mapping tokens to experts.
        when moe_config.num_experts_chosen = 1, use top1 routing;
        when moe_config.num_experts_chosen > 1, use topk routing

        Args:
            d_model (int): The hidden size of each token.
            moe_config(MoEConfig): The configuration of MoE (Mixture of Expert).
            routing_policy: The policy of mapping tokens to experts. Default: topkRouter
            training (bool): The value indicating whether is in training phase.
            parallel_config: The parallel-related configuration.
        Inputs:
            - **input_tensor** (Tensor) - Tensor of shape :math:`(expert\_parallel, tokens\_per\_device,
            hidden\_size)`.

        Outputs:
            Tensor of shape :math:`(expert\_parallel, tokens\_per\_device, expert\_dim)`.
    """

    def __init__(self,
                 d_model,
                 moe_config,
                 routing_policy=None,
                 training=True,
                 parallel_config=None,
                 init_method_std=0.01):
        super(Router, self).__init__()
        dp = parallel_config.data_parallel * parallel_config.context_parallel
        self.d_model = d_model
        self.moe_config = moe_config
        self.expert_dim = moe_config.expert_num
        self.capacity_factor = moe_config.capacity_factor
        self.num_experts_chosen = moe_config.num_experts_chosen
        self.training = training
        self.routing_policy = routing_policy
        self.noisy_policy = None  # candidate: ["jitter", "rsample", "None"]
        self.noisy_epsilon = 1e-2
        self.noise = Tensor(np.random.uniform(1 - self.noisy_epsilon, 1 + self.noisy_epsilon, (d_model,)))
        weight_init = Normal(sigma=init_method_std, mean=0)
        self.dense = Dense(in_channels=self.d_model, out_channels=self.expert_dim, weight_init=weight_init,
                           has_bias=False, dtype=moe_config.router_dense_type)
        self.dense.matmul.shard(((dp, 1), (1, 1)))
        self.mul = P.Mul()
        self.cast = P.Cast()

        if self.routing_policy == "TopkRouterV1":
            self.router = TopkRouter(d_model=d_model, moe_config=moe_config, training=training,
                                     parallel_config=parallel_config)
        elif self.routing_policy == "TopkRouterV2":
            self.router = TopkRouterV2(d_model=d_model, moe_config=moe_config, training=training,
                                       parallel_config=parallel_config)
        else:
            self.router = routing_policy

        if not (_get_parallel_mode() in (ParallelMode.AUTO_PARALLEL,) and _is_sharding_propagation()):
            self.mul.shard(((dp, 1, 1), (dp,)))

    def construct(self, input_tensor, seq_chunk=None):
        """ Router construct"""
        input_tensor = self.cast(input_tensor, self.moe_config.router_dense_type)
        if self.noisy_policy == "jitter" and self.training:
            # Here, we temporarily implement the multiplicative jitter this way,
            # for the lack of UniforReal parallel operator.
            input_tensor = self.mul(input_tensor, self.noise)

        router_logits = self.dense(input_tensor)
        if self.routing_policy == "TopkRouterV2":
            return self.router(router_logits, seq_chunk=seq_chunk)
        return self.router(router_logits)


class TopkRouter(Cell):
    r"""
        A router implementation which maps each tokens to the topk expert.

        Args:
            d_model (int): The hidden size of each token.
            moe_config(MoEConfig): The configuration of MoE (Mixture of Expert).
            training (bool): The value indicating whether is in training phase.
            config: The parallel-related configuration.
        Inputs:
            - **input_tensor** (Tensor) - Tensor of shape :math:`(expert\_parallel, tokens\_per\_device,
            hidden\_size)`.

        Outputs:
            Tensor of shape :math:`(expert\_parallel, tokens\_per\_device, expert\_dim, expert\_capacity)`,
            Tensor of shape :math:`(expert\_parallel, tokens\_per\_device, expert\_dim, expert\_capacity)`,
            Tensor of shape :math:`(1)`.
    """

    def __init__(self,
                 d_model,
                 moe_config,
                 training=True,
                 parallel_config=None):
        super(TopkRouter, self).__init__()
        if _get_parallel_mode() in (ParallelMode.AUTO_PARALLEL,) and _is_sharding_propagation():
            dp = parallel_config.data_parallel
            self.d_model = d_model
            self.expert_dim = moe_config.expert_num
            self.capacity_factor = moe_config.capacity_factor
            self.training = training
            self.dp_group = dp
            self.noisy_policy = None
            self.cast = P.Cast()
            self.reshape = P.Reshape()
            self.shape = P.Shape()
            self.softmax = P.Softmax(axis=-1)
            self.argmax = P.ArgMaxWithValue(axis=-1, keep_dims=False)
            self.num_experts_chosen = moe_config.num_experts_chosen
            self.onehot = P.OneHot()
            self.onehot2 = P.OneHot()
            self.onehot3 = P.OneHot()
            self.on_value = Tensor(1.0, mstype.float32)
            self.off_value = Tensor(0.0, mstype.float32)

            self.reduce_mean = P.ReduceMean(keep_dims=False)
            self.reduce_mean2 = P.ReduceMean(keep_dims=False)
            self.reduce_mean3 = P.ReduceMean(keep_dims=False)
            self.mul = P.Mul()
            self.mul2 = P.Mul()
            self.mul3 = P.Mul()
            self.mul4 = P.Mul()
            self.mul5 = P.Mul()
            self.mul6 = P.Mul()
            self.mul7 = P.Mul()
            self.mul8 = P.Mul().shard(((dp, 1, 1), (dp, 1, 1)))
            self.mul9 = P.Mul().shard(((dp, 1, 1, 1), (dp, 1, 1, 1)))
            self.not_equal = P.NotEqual()
            self.div1 = P.RealDiv()
            self.div2 = P.RealDiv()
            self.add = P.Add()
            self.add1 = P.Add()
            self.add2 = P.Add()
            self.add3 = P.Add()
            self.add4 = P.Add()
            self.sub = P.Sub()

            self.cumsum = P.CumSum(exclusive=True)
            self.less = P.Less()
            self.reduce_sum = P.ReduceSum(keep_dims=False)
            self.reduce_sum_keep = P.ReduceSum(keep_dims=True)
            self.reduce_sum_keep2 = P.ReduceSum(keep_dims=True)
            self.expand = P.ExpandDims()
            self.expand2 = P.ExpandDims()
            self.add_scala = P.Add()
            self.init_loss = Tensor(0.0, mstype.float32)
        else:
            dp = parallel_config.data_parallel
            self.d_model = d_model
            self.expert_dim = moe_config.expert_num
            self.capacity_factor = moe_config.capacity_factor
            self.save_token_distribution = moe_config.save_token_distribution
            self.enable_cold_hot_expert = moe_config.enable_cold_hot_expert
            self.training = training
            self.dp_group = dp
            self.noisy_policy = None
            self.cast = P.Cast()
            self.reshape = P.Reshape()
            self.shape = P.Shape()
            self.softmax = P.Softmax(axis=-1).shard(((dp, 1, 1,),))
            self.argmax = P.ArgMaxWithValue(axis=-1, keep_dims=False).shard(((dp, 1, 1),))
            self.num_experts_chosen = moe_config.num_experts_chosen
            self.onehot = P.OneHot().shard(((dp, 1, 1), (), ()))
            self.onehot2 = P.OneHot().shard(((dp, 1, 1), (), ()))
            self.onehot3 = P.OneHot().shard(((dp, 1, 1, 1), (), ()))
            self.on_value = Tensor(1.0, mstype.float32)
            self.off_value = Tensor(0.0, mstype.float32)

            self.reduce_mean = P.ReduceMean(keep_dims=False).shard(((dp, 1, 1),))
            self.reduce_mean2 = P.ReduceMean(keep_dims=False).shard(((dp, 1, 1),))
            self.reduce_mean3 = P.ReduceMean(keep_dims=False).shard(((dp, 1),))
            self.mul = P.Mul().shard(((dp, 1), (dp, 1)))
            self.mul2 = P.Mul().shard(((), ()))
            self.mul3 = P.Mul().shard(((), ()))
            self.mul4 = P.Mul().shard(((dp, 1, 1), (dp, 1, 1)))
            self.mul5 = P.Mul().shard(((dp, 1, 1), (dp, 1, 1)))
            self.mul6 = P.Mul().shard(((dp, 1), (dp, 1)))
            self.mul7 = P.Mul().shard(((dp, 1), (dp, 1)))
            self.mul8 = P.Mul().shard(((dp, 1, 1), (dp, 1, 1)))
            self.mul9 = P.Mul().shard(((dp, 1, 1, 1), (dp, 1, 1, 1)))
            self.not_equal = P.NotEqual().shard(((dp, 1, 1, 1), ()))
            self.div1 = P.RealDiv().shard(((dp, 1, 1), (dp, 1, 1)))
            self.div2 = P.RealDiv().shard(((dp, 1, 1, 1), (dp, 1, 1, 1)))
            self.add = P.Add().shard(((dp, 1, 1), (dp, 1, 1)))
            self.add1 = P.Add().shard(((dp, 1, 1), ()))
            self.add2 = P.Add().shard(((dp, 1, 1, 1), (dp, 1, 1, 1)))
            self.add3 = P.Add().shard(((dp, 1), (dp, 1)))
            self.add4 = P.Add().shard(((dp, 1, 1, 1), ()))
            self.sub = P.Sub().shard(((), (dp, 1, 1)))

            self.cumsum = P.CumSum(exclusive=True).shard(((dp, 1, 1),))
            self.less = P.Less().shard(((dp, 1, 1), ()))
            self.reduce_sum = P.ReduceSum(keep_dims=False).shard(((dp, 1, 1),))
            self.reduce_sum_keep = P.ReduceSum(keep_dims=True).shard(((dp, 1, 1),))
            self.reduce_sum_keep2 = P.ReduceSum(keep_dims=True).shard(((dp, 1, 1, 1),))
            self.expand = P.ExpandDims().shard(((dp, 1),))
            self.expand2 = P.ExpandDims().shard(((dp, 1, 1),))
            self.add_scala = P.Add().shard(((), ()))
            self.init_loss = Tensor(0.0, mstype.float32)
            if self.save_token_distribution:
                self.cur_layer = moe_config.cur_layer
                self.tensor_summary = P.TensorSummary()
            if self.enable_cold_hot_expert:
                self.cur_layer = moe_config.cur_layer
                self.cumsum_value = Parameter(initializer('zeros', (self.expert_dim,), mstype.int32),
                                              name="cumsum_value" + str(self.cur_layer), requires_grad=False,
                                              parallel_optimizer=False)
                self.assign = P.Assign().shard(((1,), (1,)))

    def construct(self, router_logits):
        """forward process"""
        router_logits_shape = self.shape(router_logits)
        router_logits = self.reshape(router_logits, (-1, router_logits_shape[-1]))
        logits_shape = self.shape(router_logits)
        tokens_per_group = logits_shape[0] // self.dp_group
        expert_capacity = calculate_expert_capacity(self.num_experts_chosen, tokens_per_group, self.capacity_factor,
                                                    self.expert_dim)
        router_logits = self.reshape(router_logits, (self.dp_group, tokens_per_group, self.expert_dim))

        accum_expert_mask = 0
        accum_expert_gate = 0
        loss = self.init_loss
        mask_count = 0
        accum_combine_tensor = 0
        # Probabilities for each token of what expert is should be sent to
        router_prob = self.softmax(router_logits)

        for expert_chosen_index in range(self.num_experts_chosen):
            # for each token, set the router_prob of the selected experts to zero
            router_prob = self.mul4(router_prob, self.sub(self.on_value, accum_expert_mask))
            # shape is : (dp_group, tokens_per_group)
            expert_index, expert_gate = self.argmax(router_prob)
            # expert_mask's shape: (dp_group, tokens_per_group, self.expert_dim)
            expert_mask = self.onehot(expert_index, self.expert_dim, self.on_value, self.off_value)
            # renormalize the rest prob to be of sum 1
            router_prob_normal = self.div1(router_prob, self.add1(self.reduce_sum_keep(router_prob, -1), 1e-9))

            # the balance loss is computed at each routing step
            loss = self.add_scala(loss, self._auxiliary_loss(expert_mask, router_prob_normal))

            output = self._maskout_overflowed_tokens(expert_mask, expert_capacity, expert_gate,
                                                     mask_count, expert_chosen_index)
            expert_mask, expert_gate, expert_mask_flat, position_in_expert = output[0], output[1], output[2], output[3]
            accum_expert_mask = self.add(accum_expert_mask, expert_mask)
            accum_expert_gate = self.add3(accum_expert_gate, expert_gate)
            mask_count = self.add(mask_count, self.reduce_sum_keep(expert_mask, 1))

            # combine_tensor's shape: (dp_group, tokens_per_group)
            combine_tensor = self.mul7(expert_gate, expert_mask_flat)
            # combine_tensor's shape: (dp_group, tokens_per_group, self.expert_dim)
            combine_tensor = self.mul8(self.expand(combine_tensor, -1),
                                       self.onehot2(expert_index, self.expert_dim, self.on_value, self.off_value))
            # combine_tensor's shape: (dp_group, tokens_per_group, self.expert_dim, self.expert_capacity)
            combine_tensor = self.mul9(self.expand2(combine_tensor, -1),
                                       self.onehot3(self.cast(position_in_expert, mstype.int32), expert_capacity,
                                                    self.on_value, self.off_value))
            accum_combine_tensor = self.add2(accum_combine_tensor, combine_tensor)

        # expert weights normalization
        combine_tensor_sum = self.reduce_sum_keep2(self.reduce_sum_keep2(accum_combine_tensor, -1), -2)
        accum_combine_tensor = self.div2(accum_combine_tensor, self.add4(combine_tensor_sum, 1e-9))
        # dispatch_tensor is of boolean type. Here, using NotEqual instead of Cast, for that 'Cast to bool' has
        # bad performance
        dispatch_tensor = self.not_equal(accum_combine_tensor, 0.0)
        return dispatch_tensor, accum_combine_tensor, loss

    def _auxiliary_loss(self, expert_mask, router_prob):
        """
        Computing the load balance loss.
        """
        # density_1's shape: (dp_group, self.expert_dim)
        density_1 = self.reduce_mean(expert_mask, 1)
        # density_1_proxy's shape: (dp_group, self.expert_dim)
        density_1_proxy = self.reduce_mean2(router_prob, 1)
        loss = self.mul(density_1, density_1_proxy)
        loss = self.reduce_mean3(loss)
        loss = self.mul3(self.mul2(loss, self.expert_dim), self.expert_dim)
        return loss

    def _maskout_overflowed_tokens(self, expert_mask, expert_capacity, expert_gate, last_num, expert_chosen_index):
        """
        Keeping only the tokens that fit within expert_capacity.
        """
        cumsum = self.cumsum(expert_mask, 1)
        if expert_chosen_index > 0:
            cumsum = self.add(cumsum, last_num)
        if self.save_token_distribution:
            record_name = 'layer-' + str(self.cur_layer)
            self.tensor_summary(record_name, cumsum[0][-1])
        if self.enable_cold_hot_expert:
            cumsum_int_value = self.cast(cumsum[0][-1], mstype.int32)
            self.assign(self.cumsum_value, cumsum_int_value)
        # position_in_expert's shape: (dp_group, tokens_per_group, self.expert_dim)
        position_in_expert = self.mul4(cumsum, expert_mask)
        less_result = self.less(position_in_expert, expert_capacity)
        # expert_mask's shape: (dp_group, tokens_per_group, self.expert_dim)
        expert_mask = self.mul5(less_result, expert_mask)
        # expert_mask_flat's shape: (dp_group, tokens_per_group)
        expert_mask_flat = self.reduce_sum(expert_mask, -1)

        # Mask out the experts that have overflowed the expert_capacity.
        # expert_gate's shape: (dp_group, tokens_per_group)
        expert_gate = self.mul6(expert_gate, expert_mask_flat)
        output = (expert_mask, expert_gate, expert_mask_flat, position_in_expert)
        return output


# pylint: disable=C0330
class TopkRouterV2(Cell):
    r"""
        A router implementation which maps each tokens to the topk expert.

        Args:
            d_model (int): The hidden size of each token.
            moe_config(MoEConfig): The configuration of MoE (Mixture of Expert).
            training (bool): The value indicating whether is in training phase.
            config: The parallel-related configuration.
        Inputs:
            - **router_logits** (Tensor) - Tensor of shape :math:`(data\_parallel, tokens\_per\_group,
            expert\_dim)`.(dp, N, expert_dim)

        Outputs:
            - **dispatch_index** (Tensor) - Tensor of shape :math:`(data\_parallel, expert\_dim, expert\_capacity)`,
            - **combine_index** (Tensor) - Tensor of shape :math:`(data\_parallel, tokens\_per\_group, k)`,
            - **router_coeff** (Tensor) - Tensor of shape :math:`(data\_parallel, tokens\_per\_group, k)`.
    """

    def __init__(self,
                 d_model,
                 moe_config,
                 training=True,
                 parallel_config=None):
        super(TopkRouterV2, self).__init__()

        dp = parallel_config.data_parallel * parallel_config.context_parallel
        self.mp = parallel_config.model_parallel
        self.ep = parallel_config.expert_parallel
        self.d_model = d_model
        self.moe_config = moe_config
        self.expert_dim = moe_config.expert_num
        self.egroup_size = moe_config.expert_group_size
        self.capacity_factor = moe_config.capacity_factor
        self.save_token_distribution = moe_config.save_token_distribution
        self.training = training
        self.dp_group = dp
        self.num_experts_chosen = moe_config.num_experts_chosen
        self.on_value = Tensor(1.0, mstype.float32)
        self.off_value = Tensor(0.0, mstype.float32)
        self.range = Tensor(np.tile(np.arange(moe_config.max_router_load) + 1,
                                    (self.num_experts_chosen, 1)), mstype.float32)
        # aux loss free
        if self.moe_config.balance_via_topk_bias:
            self.topk_bias = Parameter(initializer('zeros', (self.expert_dim), mstype.float32),
                                       requires_grad=False, parallel_optimizer=False)
            self.one_over_expert_dim = Tensor([1 / self.expert_dim], mstype.float32)
            self.sign = P.Sign().shard(((1,),))
            self.gate_gather = P.Gather(batch_dims=2).shard(((dp, 1, 1), (dp, 1, 1)))
            self.afb_add = P.Add().shard(((1,), (1,)))
            self.afb_sub = P.Sub().shard(((1,), (1,)))
            self.afb_add_topk_bias = P.Add().shard(((dp, 1, 1), (1,)))
            self.afb_add_topk_bias.recompute(False)
            self.assign = P.Assign().shard(((1,), (1,)))
            self.assign.recompute(False)
            self.afb_mul = P.Mul().shard(((1,), ()))
            self.afb_reduce_mean = P.ReduceMean(keep_dims=False).shard(((1, 1),))
            self.afb_topk = P.TopK().shard(((dp, 1, 1),))
            self.afb_topk.recompute(False)
            self.expert_load = Parameter(initializer('zeros', (self.expert_dim), mstype.float32),
                                         requires_grad=False, parallel_optimizer=False)
            self.assign_add = P.AssignAdd().shard(((1,), (1,)))
            self.assign_add.recompute(False)

        if self.moe_config.topk_method == "noaux_tc":
            self.tc_topk = P.TopK().shard(((dp, 1, 1, 1),))
            self.tc_topk.recompute(False)
            self.tc_sum = P.ReduceSum(keep_dims=False).shard(((dp, 1, 1, 1),))

        self.cast = P.Cast()
        self.reshape = P.Reshape()
        self.shape = P.Shape()
        self.gating_activation = P.Softmax(axis=-1).shard(((dp, 1, 1,),)) \
            if not moe_config.use_gating_sigmoid else P.Sigmoid().shard(((dp, 1, 1,),))
        self.topk = P.TopK().shard(((dp, 1, 1),))
        self.topk.recompute(False)
        self.argmax = P.ArgMaxWithValue(axis=-1, keep_dims=False).shard(((dp, 1, 1),))
        self.onehot_2d = P.OneHot().shard(((dp, 1, 1), (), ()))
        self.onehot_3d = P.OneHot().shard(((dp, 1, 1, 1), (), ()))
        self.cumsum = P.CumSum(exclusive=False).shard(((dp, 1, 1),))
        self.mul_2d_1d = P.Mul().shard(((dp, 1), ()))
        self.mul_2d = P.Mul().shard(((dp, 1), (dp, 1)))
        self.mul_3d = P.Mul().shard(((dp, 1, 1), (dp, 1, 1)))
        self.mul_4d = P.Mul().shard(((dp, 1, 1, 1), (dp, 1, 1, 1)))
        self.add_2d = P.Add().shard(((dp, 1), (dp, 1)))
        self.add_3d = P.Add().shard(((dp, 1, 1), (dp, 1, 1)))
        self.less = P.Less().shard(((dp, 1, 1), ()))
        self.gt = P.Greater().shard(((dp, 1), ()))
        self.reduce_sum = P.ReduceSum(keep_dims=False).shard(((dp, 1, 1),))
        self.transpose_3d = P.Transpose().shard(((dp, 1, 1),))
        self.transpose = P.Transpose().shard(((dp, 1, 1, 1),))
        self.slice = P.StridedSlice().shard(((dp, 1, 1),))
        self.slice_range = P.StridedSlice().shard(((1, 1),))
        self.bmm_range = P.BatchMatMul().shard(((1, 1), (dp, 1, 1, 1)))
        self.add_eps = P.Add().shard(((dp, 1, 1), ()))
        self.reduce_sum_keep = P.ReduceSum(keep_dims=True).shard(((dp, 1, 1),))
        self.div_3d = P.RealDiv().shard(((dp, 1, 1), (dp, 1, 1)))
        self.concat_3d = P.Concat(1).shard(((dp, 1, 1), (dp, 1, 1)))
        self.zeros_op = P.Zeros().shard(((dp, 1, 1),))
        self.not_equal = P.NotEqual().shard(((dp, 1, 1), ()))

        # sort indexing
        self.range2 = Tensor(np.tile(np.arange(moe_config.max_router_load),
                                     (self.expert_dim, 1)), mstype.float32)
        self.add_one = P.Add().shard(((dp, 1, 1), ()))
        self.add_range = P.Add().shard(((1, 1, 1), ()))
        self.sub_range = P.Sub().shard(((), (dp, 1, 1)))
        self.mul_range = P.Mul().shard(((dp, 1, 1), (1, 1, 1)))
        self.sort_range = P.Sort().shard(((dp, 1, 1),))
        self.mod = P.Mod().shard(((dp, 1, 1), ()))
        self.print = P.Print()

        # group limited greedy
        self.tile = P.Tile()

        # auxiliary loss
        self.z_loss_func = ZLoss(parallel_config=parallel_config)
        self.z_loss_coeff = Tensor(moe_config.z_loss_factor, mstype.float32)
        # for auxiliary balancing loss
        self.mul = P.Mul().shard(((), ()))
        self.reduce_mean = P.ReduceMean(keep_dims=False).shard(((dp, 1, 1),))
        self.reduce_mean_2d = P.ReduceMean(keep_dims=False).shard(((dp, 1),))
        self.add_scalar = P.Add()
        # auxiliary loss config
        self.aux_loss_config = dict(zip(moe_config.aux_loss_types, moe_config.aux_loss_factors))

        # dynamic capacity
        self.on_value_int = Tensor(1, mstype.int32)
        self.off_value_int = Tensor(0, mstype.int32)
        self.mod = P.Mod().shard(((dp, 1, 1), ()))
        self.add = P.Add()
        self.sub = P.Sub()
        self.mod_expert = P.Mod()
        self.tensor2scalar = TensorToScalar()
        if moe_config.callback_moe_droprate:
            self.fi_parameter = Parameter(initializer("zeros", (dp, self.expert_dim), mstype.float32),
                                        requires_grad=False, parallel_optimizer=False)
            self.assign_fi = P.Assign().shard(((dp, 1), (dp, 1)))
            self.assign_fi.recompute(False)

        # allgather dispatcer
        self.use_allgather_dispatcher = moe_config.use_allgather_dispatcher
        self.outer_dp = dp // self.ep
        self.inner_dp = self.ep
        if self.use_allgather_dispatcher:
            self.dispatch_gather = P.Gather(batch_dims=2).shard(((self.outer_dp, 1, 1, 1), (self.outer_dp, 1, 1, 1),))
            self.reshape_redist = P.Reshape().add_prim_attr("skip_redistribution", True)
            self.tile_5d = P.Tile().shard(((self.outer_dp, 1, self.ep, 1, 1),))
            self.concat = P.Concat(3).shard(((self.outer_dp, 1, self.ep, 1, 1), (self.outer_dp, 1, self.ep, 1, 1)))
            self.zeros = Tensor(
                np.zeros((self.outer_dp, self.inner_dp, self.ep * self.expert_dim, 1, d_model)), mstype.float16)
            self.tile_idx = P.Tile().shard(((self.outer_dp, self.inner_dp, 1, 1, 1),))
            self.tile_coeff = P.Tile().shard(((self.outer_dp, self.inner_dp, 1, 1, 1, 1),))
            self.combine_gather = P.Gather(batch_dims=3).shard(((self.outer_dp, self.inner_dp, 1, 1, 1),
                                                                (self.outer_dp, self.inner_dp, 1, 1, 1)))
            self.mul_router_coeff = P.Mul().shard(((self.outer_dp, self.inner_dp, 1, 1, 1, 1),
                                                   (self.outer_dp, self.inner_dp, 1, 1, 1, 1)))
            self.sum_router_coeff = P.ReduceSum(keep_dims=False).shard(((self.outer_dp, self.inner_dp, 1, 1, 1, 1),))
            self.reduce_sum_5d = P.ReduceSum(keep_dims=False).shard(((self.outer_dp, 1, 1, 1, 1),))
            ################# help for mask #############
            ones_per_row = self.expert_dim // self.ep
            indices = np.arange(ones_per_row) + (np.arange(dp) * ones_per_row)[:, None] % self.expert_dim
            one_hot_array = np.zeros((dp, self.expert_dim), dtype=int)
            one_hot_array[np.arange(dp)[:, None], indices] = 1
            self.one_hot = Tensor(np.tile(one_hot_array,
                                          (self.ep, 1)
                                          ).reshape(self.outer_dp, self.inner_dp, self.ep, self.expert_dim, 1),
                                  mstype.float16)
            self.mul_onehot = P.Mul().shard(((self.outer_dp, 1, self.ep, 1, 1), (self.outer_dp, 1, self.ep, 1, 1)))
            ################# help for mask end ##########
        else:
            self.dispatch_gather = P.Gather(batch_dims=1).shard(((dp, 1, 1), (dp, 1, 1),))
            self.concat = P.Concat(2).shard(((dp, 1, 1, 1), (dp, 1, 1, 1)))
            self.zeros_op_4d = P.Zeros().shard(((dp, 1, 1, 1),))
            self.combine_gather = P.Gather(batch_dims=1).shard(((dp, 1, 1), (dp, 1, 1),))
            self.mul_router_coeff = P.Mul().shard(((dp, 1, 1, 1), (dp, 1, 1, 1)))
            self.sum_router_coeff = P.ReduceSum(keep_dims=False).shard(((dp, 1, 1, 1),))

        self.seq_split_num = parallel_config.seq_split_num
        self.seq_pipe = self.seq_split_num > 1
        if self.seq_pipe:
            self.pi_cache = Parameter(initializer('zeros', (dp, self.expert_dim), mstype.float32),
                                      requires_grad=False, parallel_optimizer=False)
            self.fi_cache = Parameter(initializer('zeros', (dp, self.expert_dim), mstype.float32),
                                      requires_grad=False, parallel_optimizer=False)
            self.assign_cache = P.Assign().shard(((dp, 1), (dp, 1)))
            self.add_cache = P.Add().shard(((dp, 1), (dp, 1)))
            self.mul_cache = P.Mul().shard(((dp, 1), ()))
        # topkrouter
        if self.moe_config.use_fused_ops_topkrouter:
            # pylint: disable=W0212
            self.topkrouter = P._inner_ops.TopKRouter().shard(((dp, 1, 1),))
        # interleave
        if self.moe_config.comp_comm_parallel:
            self.comp_comm_parallel_degree = self.moe_config.comp_comm_parallel_degree
        else:
            self.comp_comm_parallel_degree = 1

    def dispatch(self, input_tensor, dispatch_index):
        r"""
            Implementing dispatch operation.
            Inputs:
                - **input_tensor** (Tensor) - Tensor of shape :math:`(data\_parallel, tokens\_per\_group,
                hidden\_size)`.(dp, N, h),
                - **dispatch_index** (Tensor) - Tensor of shape :math:`(data\_parallel, expert\_num,
                expert\_capacity)`.(dp, E, n).

            Outputs:
                - **expert_input** (Tensor) - Tensor of shape :math:`(data\_parallel, expert\_num,
                expert\_capacity, hidden\_size)`.(dp, E, n, h).
        """
        zeros_3d = self.zeros_op((input_tensor.shape[0], 1, input_tensor.shape[-1]),
                                 dtype=F.dtype(input_tensor))
        input_tensor_padded = self.concat_3d((zeros_3d, input_tensor)) # (dp, 1+N, h) <-- (dp, N, h)
        if self.use_allgather_dispatcher:
            # input_tensor_padded shape: (outer_dp, inner_dp, 1+N, h) <-- (dp, 1+N, h)
            input_tensor_padded = self.reshape(input_tensor_padded, (
                self.outer_dp, self.inner_dp, input_tensor_padded.shape[1], input_tensor_padded.shape[2]
            ))
            # dispatch_index shape: (outer_dp, inner_dp, E, n) <-- (dp, E, n)
            dispatch_index = self.reshape(dispatch_index, (
                self.outer_dp, self.inner_dp, dispatch_index.shape[1], dispatch_index.shape[2]))
            # expert_input shape: (outer_dp, innert_dp, E, n, h)
            expert_input = self.dispatch_gather(input_tensor_padded, dispatch_index, 2)
            # expert_input shape: (dp, E, n, h)
            expert_input = self.reshape(expert_input, (
                self.dp_group, self.expert_dim, expert_input.shape[-2], expert_input.shape[-1]))
        else:
            # expert_input shape: (dp, E, n, h) <-- (dp, N, h), (dp, E, n)
            expert_input = self.dispatch_gather(input_tensor_padded, dispatch_index, 1)
        return expert_input

    def combine(self, expert_output, combine_index, router_coeff):
        r"""
            Implementing combine operation.
            Inputs:
                - **expert_output** (Tensor) - Tensor of shape :math:`(data\_parallel, expert\_num,
                expert\_capacity, hidden\_size)`.(dp, E, n, h),
                - **combine_index** (Tensor) - Tensor of shape :math:`(data\_parallel, tokens\_per\_group,
                num\_experts\_chosen)`.(dp, N, k),
                - **router_coeff** (Tensor) - Tensor of shape :math:`(data\_parallel, tokens\_per\_group,
                num\_experts\_chosen)`.(dp, N, k).

            Outputs:
                - **output_tensor** (Tensor) - Tensor of shape :math:`(data\_parallel, tokens\_per\_group,
                hidden\_size)`.(dp, N, h).
        """
        if self.use_allgather_dispatcher:
            hidden_size = expert_output.shape[-1]
            # expert_output shape: (outer_dp, inner_dp, ep, E//ep, n*h)
            expert_output = self.reshape(expert_output, (
                self.outer_dp, self.inner_dp, self.ep, self.expert_dim // self.ep, -1))
            # expert_output shape: (outer_dp, inner_dp, ep, E, n*h)
            expert_output = self.tile_5d(expert_output, (1, 1, 1, self.ep, 1))
            # one_hot shape: (outer_dp, inner_dp, ep, E, 1)
            # expert_output shape: (outer_dp, inner_dp, ep, E, n*h)
            expert_output = self.mul_onehot(self.one_hot, expert_output)
            # expert_output shape: (outer_dp, inner_dp, ep*E, n, h)
            expert_output = self.reshape(expert_output, (
                self.outer_dp, self.inner_dp, self.expert_dim * self.ep, -1, hidden_size
            ))
            # expert_output shape: (outer_dp, inner_dp, ep*E, 1+n, h)
            expert_output = self.concat(
                (self.cast(self.zeros, F.dtype(expert_output)), expert_output))
            # expert_output shape: (outer_dp, inner_dp, ep, E*(1+n), h)
            expert_output = self.reshape_redist(expert_output, (
                self.outer_dp, self.inner_dp, self.ep, -1, hidden_size))
            # combine_index shape: (outer_dp, 1, ep, n, k) <-- (dp, n, k)
            combine_index = self.reshape(combine_index, (
                self.outer_dp, 1, self.ep, -1, combine_index.shape[2]))
            # combine_index shape: (outer_dp, inner_dp, ep, n, k)
            combine_index = self.tile_idx(combine_index, (1, self.inner_dp, 1, 1, 1))
            # output_tensor shape: (outer_dp, inner_dp, ep, n, k, h)
            output_tensor = self.combine_gather(expert_output, combine_index, 3)
            router_coeff = self.cast(router_coeff, F.dtype(expert_output))
            # router_coeff shape: (outer_dp, 1, ep, n, k, 1)
            router_coeff = self.reshape(router_coeff, (
                self.outer_dp, 1, self.ep, router_coeff.shape[1], router_coeff.shape[2], 1))
            # router_coeff shape: (outer_dp, inner_dp, ep, n, k, 1)
            router_coeff = self.tile_coeff(router_coeff, (
                1, self.inner_dp, 1, 1, 1, 1))
            # output_tensor shape: (outer_dp, inner_dp, ep, n, k, h)
            output_tensor = self.mul_router_coeff(output_tensor, router_coeff)
            # output_tensor shape: (outer_dp, inner_dp, ep, n, h) <-- (outer_dp, inner_dp, ep, n, k, h)
            output_tensor = self.sum_router_coeff(output_tensor, 4)
            # output_tensor shape: (outer_dp, ep, n, h) <-- (outer_dp, inner_dp, ep, n, h)
            output_tensor = self.reduce_sum_5d(output_tensor, 1)
            # output_tensor shape: (dp, n, h) <-- (outer_dp, ep, n, h)
            output_tensor = self.reshape(output_tensor, (
                self.dp_group, output_tensor.shape[2], hidden_size))
        else:
            zeros = self.zeros_op_4d((expert_output.shape[0],
                                      expert_output.shape[1],
                                      1,
                                      expert_output.shape[-1]), dtype=F.dtype(expert_output))
            expert_output = self.concat(
                (zeros, expert_output))  # (dp, E, 1+n, h) <-- (dp, E, n, h)
            expert_output = self.reshape(
                expert_output, (expert_output.shape[0],
                                expert_output.shape[1] * expert_output.shape[2],
                                expert_output.shape[3]))  # (dp, E*(1+n), h) <-- (dp, E, 1+n, h)
            output_tensor = self.combine_gather(
                expert_output, combine_index, 1)  # (dp, n, k, h) <-- (dp, E*(1+n), h), (dp, n, k)
            router_coeff = self.cast(router_coeff, F.dtype(expert_output))
            # output_tensor shape: (dp, n, k, h) <-- (dp, n, k, h) (dp, n, k, 1)
            output_tensor = self.mul_router_coeff(
                output_tensor,
                self.reshape(router_coeff, (router_coeff.shape[0], router_coeff.shape[1], router_coeff.shape[2], 1)))
            # output_tensor shape: (dp, n, h) <-- (dp, n, k, h)
            output_tensor = self.sum_router_coeff(output_tensor, 2)
        return output_tensor

    def construct(self, router_logits, seq_chunk=None):
        """
        Calculate dispatch_policy, combine_policy, router_coeff.
        """
        z_loss = self.z_loss_func(router_logits, self.z_loss_coeff)
        extra_loss = z_loss

        router_prob = self.gating_activation(router_logits)  # (dp, N, expert_dim)fp32 <-- (dp, N, expert_dim)fp32
        # normalize if use sigmoid activate for auxiliary loss
        router_prob_for_aux = self._normalize(router_prob) if self.moe_config.use_gating_sigmoid else router_prob
        if self.moe_config.topk_method in ("group_limited_greedy", "noaux_tc"):
            if self.moe_config.balance_via_topk_bias:
                router_prob_with_bias = self.afb_add_topk_bias(router_prob, self.topk_bias)
            else:
                router_prob_with_bias = router_prob
            group_scores = self.reshape(router_prob_with_bias,
                                        (router_prob_with_bias.shape[0],
                                         router_prob_with_bias.shape[1],
                                         self.moe_config.n_group, -1))
            if self.moe_config.topk_method == "noaux_tc":
                top_2, _ = self.tc_topk(group_scores, 2)
                group_scores = self.tc_sum(top_2, -1)
            else:
                group_scores = group_scores.max(axis=-1)
            top_values, _ = self.topk(group_scores, self.moe_config.topk_group)  # (dp, N, top_k)
            group_mask = self.cast(
                group_scores.ge(self.tile(top_values.min(axis=-1, keepdims=True),
                                          (1, 1, self.moe_config.n_group))), mstype.float32)
            score_mask = self.reshape(
                group_mask, (group_mask.shape[0], group_mask.shape[1], group_mask.shape[2], 1))  # (dp, N,  n_group, 1)
            score_mask = score_mask.repeat_interleave(self.moe_config.expert_num // self.moe_config.n_group, dim=-1)
            score_mask = self.reshape(
                score_mask, (score_mask.shape[0], score_mask.shape[1], -1))  # (dp, N, n_routed_experts)
            tmp_scores = ops.masked_fill(router_prob_with_bias, ~score_mask.bool(), 0.0)

            if self.moe_config.balance_via_topk_bias:
                _, expert_index = self.afb_topk(tmp_scores, self.num_experts_chosen)
                expert_gate = self.gate_gather(router_prob, expert_index, 2)
                self._update_expert_load(expert_index)
            else:
                expert_gate, expert_index = self.topk(tmp_scores, self.num_experts_chosen)
        else:
            # in default, normal topk will be used
            if self.moe_config.balance_via_topk_bias:
                _, expert_index = self.afb_topk(self.afb_add_topk_bias(router_prob, self.topk_bias),
                                                self.num_experts_chosen)
                expert_gate = self.gate_gather(router_prob, expert_index, 2)
                self._update_expert_load(expert_index)
            else:
                expert_gate, expert_index = self.topk(router_prob, self.num_experts_chosen)

        if self.num_experts_chosen > 1 and self.moe_config.norm_topk_prob:
            expert_gate = self._normalize(expert_gate)  # (dp, N, k) <-- (dp, N, k)

        expert_gate = ops.mul(self.moe_config.routed_scaling_factor, expert_gate)  # (dp, N, k) <-- (dp, N, k)
        aux_loss_config_is_valid = self.aux_loss_config.get("expert", 0) > 0 \
            or self.aux_loss_config.get("device", 0) > 0 or self.aux_loss_config.get("comm", 0) > 0
        if self.moe_config.balance_via_topk_bias and aux_loss_config_is_valid:
            _, expert_index_for_aux = self.topk(router_prob_for_aux, self.num_experts_chosen)
        else:
            expert_index_for_aux = expert_index
        if self.aux_loss_config.get("expert", 0):
            expert_load_loss = self._expert_load_balancing(router_prob_for_aux, expert_index_for_aux,
                                                           self.aux_loss_config.get("expert"), seq_chunk=seq_chunk)
            extra_loss = self.add_scalar(extra_loss, expert_load_loss)
        if self.aux_loss_config.get("device", 0):
            device_load_loss = self._device_load_balancing(router_prob_for_aux, expert_index_for_aux,
                                                           self.aux_loss_config.get("device"))
            extra_loss = self.add_scalar(extra_loss, device_load_loss)
        if self.aux_loss_config.get("comm", 0):
            comm_load_loss = self._comm_load_balancing(router_prob_for_aux, expert_index_for_aux,
                                                       self.aux_loss_config.get("comm"))
            extra_loss = self.add_scalar(extra_loss, comm_load_loss)

        if self.moe_config.enable_sdrop:
            if self.moe_config.use_fused_ops_topkrouter:
                # (dp, E, n)int32, (dp, N, k), (dp, N, k) <-- (dp, N, k), (dp, N, k)
                dispatch_idx, combine_idx, router_coeff = self._maskout_overflowed_tokens_use_topkrouter(expert_index,
                                                                                                         expert_gate)
            else:
                # (dp, E, n)int32, (dp, N, k), (dp, N, k) <-- (dp, N, k), (dp, N, k)
                dispatch_idx, combine_idx, router_coeff = self._maskout_overflowed_tokens_sort_sdrop(expert_index,
                                                                                                     expert_gate)
        else:
            if self.moe_config.use_fused_ops_topkrouter:
                # (dp, E, n)int32, (dp, N, k), (dp, N, k) <-- (dp, N, k), (dp, N, k)
                dispatch_idx, combine_idx, router_coeff = self._maskout_overflowed_tokens_use_topkrouter(expert_index,
                                                                                                         expert_gate)
            else:
                # (dp, E, n)int32, (dp, N, k), (dp, N, k) <-- (dp, N, k), (dp, N, k)
                dispatch_idx, combine_idx, router_coeff = self._maskout_overflowed_tokens_sort_kdrop(expert_index,
                                                                                                     expert_gate)
        return dispatch_idx, combine_idx, router_coeff, extra_loss  # (dp, E, n)int32, (dp, N, k), (dp, N, k)

    def _update_expert_load(self, expert_index):
        expert_index = self.reshape(expert_index, (expert_index.shape[0], -1))
        expert_mask = self.onehot_2d(expert_index, self.expert_dim, self.on_value, self.off_value)
        expert_load_data = self.reduce_mean(expert_mask, 1)
        expert_load_data = self.afb_reduce_mean(expert_load_data, 0)
        self.assign_add(self.expert_load, expert_load_data)

    def _maskout_overflowed_tokens_sort_kdrop(self, expert_index, expert_gate):
        """
        Keeping only the tokens that fit within expert_capacity.
        """
        k = self.num_experts_chosen
        tokens_per_group = self.shape(expert_index)[1]
        kn = k * tokens_per_group  # this n refers to N
        if self.capacity_factor > 0:
            expert_capacity = calculate_expert_capacity_v2(self.num_experts_chosen, tokens_per_group,
                                                           self.capacity_factor, self.expert_dim,
                                                           self.mp * self.comp_comm_parallel_degree)

        else:
            expert_capacity = self._calculate_expert_capacity_dynamic(expert_index)
        # calculate combine_index from cumsum
        expert_index = self.reshape(self.transpose_3d(expert_index, (0, 2, 1)),
                                    (expert_index.shape[0], -1))  # (dp, kN) <-- (dp, N, k) account for topk priority
        expert_mask = self.onehot_2d(expert_index,
                                     self.expert_dim,
                                     self.on_value,
                                     self.off_value)  # (dp, kN, E)fp32 <-- (dp, kN)int32
        position_in_expert = self.mul_3d(
            self.cumsum(expert_mask, 1), expert_mask)  # (dp, kN, E)fp16 <-- (dp, kN, E)fp32, (dp, kN, E)fp32

        # (dp, kN, E)fp32 <-- (dp, kN, E)fp32, (dp, kN, E)bool, where 0<=position_in_expert<(1+n)
        position_in_expert = self.mul_3d(position_in_expert, self.less(position_in_expert, expert_capacity + 1))
        position_in_expert_2d = self.reduce_sum(position_in_expert, -1)  # (dp, kN)fp32 <-- (dp, kN, E)fp32
        # (dp, kN)fp32 <-- (dp, kN)fp32, (dp, kN)fp32 where 0<= combine_index <E*(1+n), combine_index = expert_id *(1+n) + position_in_expert_2d
        combine_index = self.add_2d(
            self.mul_2d_1d(expert_index, expert_capacity + 1),
            position_in_expert_2d)
        combine_index = self.transpose_3d(
            self.reshape(combine_index, (combine_index.shape[0], k, tokens_per_group)),
            (0, 2, 1))  # (dp, N, k) <-- (dp, kN) account for topk priority
        within_capacity = self.cast(self.gt(position_in_expert_2d, 0), mstype.float32)  # (dp, kN)bool

        # calculate dispatch_index from position_in_expert_onehot
        safe_kn = 2 * kn  # factor=2 for safety
        range_kn = self.slice_range(self.range2, (0, 0),
                                    (self.expert_dim, kn),
                                    (1, 1)).reshape(1, self.expert_dim, kn)  # (1, E, kN) fp32 <-- (E, 131072)
        select = self.transpose_3d(expert_mask, (0, 2, 1))  # (dp, E, kN) fp32 <-- (dp, kN, E) fp32
        dispatch_index_raw = self.add_3d(
            self.mul_range(select, range_kn),
            self.mul_range(self.sub_range(1, select),
                           self.add_range(range_kn, safe_kn)))  # (dp, E, kN) <-- (dp, E, kN) fp32
        dispatch_index, _ = self.sort_range(dispatch_index_raw)  # (dp, E, k) <-- (dp, E, kN）
        dispatch_index = self.slice(dispatch_index, (0, 0, 0),
                                    (dispatch_index.shape[0], dispatch_index.shape[1], expert_capacity),
                                    (1, 1, 1))  # (dp, E, n) <-- (dp, E, kN) fp32
        is_safe = self.less(dispatch_index, safe_kn)  # (dp, E, n) bool
        dispatch_index = self.add_one(self.mod(dispatch_index, tokens_per_group), 1)  # (dp, E, n) fp32
        dispatch_index = self.mul_3d(dispatch_index, is_safe)  # (dp, E, n) fp32

        dispatch_index = self.cast(dispatch_index, mstype.int32)
        combine_index = self.cast(combine_index, mstype.int32)
        router_coeff = self.mul_3d(
            expert_gate,
            self.transpose_3d(
                self.reshape(within_capacity, (within_capacity.shape[0], k, tokens_per_group)),
                (0, 2, 1)))  # apply within_capacity (dp, N, k) <-- (dp, N, k), (dp, N, k) <--  (dp, kN)

        return dispatch_index, combine_index, router_coeff  # (dp, E, n), (dp, N, k), (dp, N, k)


    def _maskout_overflowed_tokens_sort_sdrop(self, expert_index, expert_gate):
        """
        Keeping only the tokens that fit within expert_capacity.
        """
        k = self.num_experts_chosen
        tokens_per_group = self.shape(expert_index)[1]
        kn = k * tokens_per_group  # this n refers to N
        if self.capacity_factor > 0:
            expert_capacity = calculate_expert_capacity_v2(self.num_experts_chosen, tokens_per_group,
                                                           self.capacity_factor, self.expert_dim,
                                                           self.mp * self.comp_comm_parallel_degree)

        else:
            expert_capacity = self._calculate_expert_capacity_dynamic(expert_index)
        # calculate combine_index from cumsum
        expert_index = self.reshape(expert_index,
                                    (expert_index.shape[0], -1))  # (dp, Nk) <-- (dp, N, k) account for topk priority
        expert_mask = self.onehot_2d(expert_index,
                                     self.expert_dim,
                                     self.on_value,
                                     self.off_value)  # (dp, Nk, E)fp32 <-- (dp, Nk)int32
        position_in_expert = self.mul_3d(self.cumsum(expert_mask, 1),
                                         expert_mask)  # (dp, Nk, E)fp16 <-- (dp, Nk, E)fp32, (dp, Nk, E)fp32
        # (dp, Nk, E)fp32 <-- (dp, Nk, E)fp32, (dp, Nk, E)bool, where 0<=position_in_expert<(1+n)
        position_in_expert = self.mul_3d(position_in_expert, self.less(position_in_expert, expert_capacity + 1))
        position_in_expert_2d = self.reduce_sum(position_in_expert, -1)  # (dp, Nk)fp32 <-- (dp, Nk, E)fp32
        # (dp, Nk)fp32 <-- (dp, Nk)fp32, (dp, Nk)fp32 where 0<= combine_index <E*(1+n), combine_index = expert_id *(1+n) + position_in_expert_2d
        combine_index = self.add_2d(self.mul_2d_1d(expert_index, expert_capacity + 1), position_in_expert_2d)
        combine_index = self.reshape(
            combine_index,
            (combine_index.shape[0], tokens_per_group, k))  # (dp, N, k) <-- (dp, Nk) account for topk priority
        within_capacity = self.cast(self.gt(position_in_expert_2d, 0), mstype.float32)  # (dp, Nk)bool

        # calculate dispatch_index from position_in_expert_onehot
        safe_kn = 2 * kn  # factor=2 for safety
        range_kn = self.slice_range(self.range2,
                                    (0, 0),
                                    (self.expert_dim, kn),
                                    (1, 1)).reshape(1, self.expert_dim, kn)  # (1, E, kN) fp32 <-- (E, 131072)
        select = self.transpose_3d(expert_mask, (0, 2, 1))  # (dp, E, Nk) fp32 <-- (dp, Nk, E) fp32
        dispatch_index_raw = self.add_3d(
            self.mul_range(select, range_kn),
            self.mul_range(
                self.sub_range(1, select),
                self.add_range(range_kn, safe_kn)))  # (dp, E, kN) <-- (dp, E, kN) fp32
        dispatch_index, _ = self.sort_range(dispatch_index_raw)  # (dp, E, k) <-- (dp, E, kN）
        dispatch_index = self.slice(dispatch_index,
                                    (0, 0, 0),
                                    (dispatch_index.shape[0], dispatch_index.shape[1], expert_capacity),
                                    (1, 1, 1))  # (dp, E, n) <-- (dp, E, kN) fp32
        is_safe = self.less(dispatch_index, safe_kn)  # (dp, E, n) bool
        dispatch_index = self.add_one(ops.floor_divide(dispatch_index, k), 1)  # (dp, E, n) fp32
        dispatch_index = self.mul_3d(dispatch_index, is_safe)  # (dp, E, n) fp32

        dispatch_index = self.cast(dispatch_index, mstype.int32)
        combine_index = self.cast(combine_index, mstype.int32)
        within_capacity = self.reshape(within_capacity, (within_capacity.shape[0], tokens_per_group, k))
        router_coeff = self.mul_3d(expert_gate, within_capacity)

        return dispatch_index, combine_index, router_coeff  # (dp, E, n), (dp, N, k), (dp, N, k)


    def _maskout_overflowed_tokens_use_topkrouter(self, expert_index, expert_gate):
        """
        Using TopKRouter calculate dispatch_policy and combine_policy.
        """
        if self.capacity_factor > 0:
            tokens_per_group = self.shape(expert_index)[1]
            expert_capacity = calculate_expert_capacity_v2(self.num_experts_chosen, tokens_per_group,
                                                           self.capacity_factor, self.expert_dim,
                                                           self.mp * self.comp_comm_parallel_degree)
        else:
            expert_capacity = self._calculate_expert_capacity_dynamic(expert_index)
        if self.moe_config.enable_sdrop:
            dispatch_index, combine_index = self.topkrouter(expert_index, expert_capacity, self.expert_dim)
        else:
            dispatch_index, combine_index = self.topkrouter(expert_index, expert_capacity, self.expert_dim, 1)
        within_capacity = self.mod(combine_index, expert_capacity + 1)
        within_capacity = self.not_equal(self.cast(within_capacity, mstype.int32), 0)
        router_coeff = self.mul_3d(within_capacity, expert_gate)

        return dispatch_index, combine_index, router_coeff

    def _normalize(self, router_coeff_raw):
        router_coeff_sum = self.reduce_sum_keep(router_coeff_raw, 2)  # (dp, N, 1) <-- (dp, N, k)
        router_coeff = self.div_3d(router_coeff_raw,
                                   self.add_eps(router_coeff_sum, 1e-9))  # (dp, N, k) <-- (dp, N, k) (dp, N, 1)
        return router_coeff  # (dp, N, k)

    def _calculate_expert_capacity_dynamic(self, expert_index):
        """
        Calculate dynamic capacity.
        """
        expert_index = self.reshape(expert_index, (self.dp_group, -1))
        expert_mask = self.onehot_2d(expert_index,
                                     self.expert_dim,
                                     self.on_value_int,
                                     self.off_value_int)  # (dp, kN, E) <- (dp, kN)
        expert_mask = self.reduce_sum(expert_mask, 1)  # (dp, E) <- (dp, kN, E)
        expert_capacity = expert_mask.max()  # (1, ) <- (dp, E)
        expert_capacity = self.cast(expert_capacity, mstype.int64)
        expert_capacity_scalar = self.tensor2scalar(expert_capacity)
        mp = self.mp * self.comp_comm_parallel_degree
        expert_capacity_scalar = (expert_capacity_scalar // mp + 1) * mp
        return expert_capacity_scalar

    def _expert_load_balancing(self, scores, top_indices, alpha, seq_chunk=None):
        """Expert level load balance loss, which regularizes the load from local batch data on each
        expert to be balanced.
        Please refer to DeepSeek-V2:
        A Strong, Economical, and Efficient Mixture-of-Experts Language Model, https://arxiv.org/abs/2405.04434
        """
        pi = self.reduce_mean(scores, 1)  # (dp, E)<- (dp, N, E), 1/N * \sum_t^T s_{i,t}

        top_indices = self.reshape(self.transpose_3d(top_indices, (0, 2, 1)),
                                   (top_indices.shape[0], -1))  # (dp, kN) <-- (dp, N, k) account for topk priority
        mask = self.onehot_2d(top_indices,
                              self.expert_dim,
                              self.on_value,
                              self.off_value)  # (dp, kN, E)fp32 <-- (dp, kN)int32
        fi = self.reduce_mean(mask, 1)  # (dp, E) <- (dp, kN, E), 1/(kN) * \sum_t^T 1(token t selects expert i)

        if seq_chunk is not None:
            is_clean = self.cast(P.NotEqual()(seq_chunk, 0), mstype.float32)
            is_return = self.cast(P.Equal()(seq_chunk, self.seq_split_num - 1), mstype.float32)
            clean_zeros_pi = self.mul_cache(self.pi_cache, is_clean)
            clean_zeros_fi = self.mul_cache(self.fi_cache, is_clean)
            clean_pi_state = self.assign_cache(self.pi_cache, clean_zeros_pi)
            clean_fi_state = self.assign_cache(self.fi_cache, clean_zeros_fi)
            pi = self.add_cache(F.depend(pi, clean_pi_state), self.pi_cache)
            fi = self.add_cache(F.depend(fi, clean_fi_state), self.fi_cache)
            pi_fi = self.reduce_mean_2d(self.mul_2d(pi, fi))
            pi_cache_state = self.assign_cache(self.pi_cache, F.depend(pi, pi_fi))
            fi_cache_state = self.assign_cache(self.fi_cache, F.depend(fi, pi_fi))
            expert_load_loss = self.mul(pi_fi, alpha * self.expert_dim ** 2)
            expert_load_loss = F.depend(expert_load_loss, (pi_cache_state, fi_cache_state))
            return expert_load_loss * is_return / (self.seq_split_num * self.seq_split_num)
        expert_load_loss = self.mul(self.reduce_mean_2d(self.mul_2d(pi, fi)),
                                    alpha * self.expert_dim ** 2)  # alpha*E \sum_i^E (f_i * P_i)
        if self.moe_config.callback_moe_droprate:
            self.assign_fi(self.fi_parameter, fi)
        return expert_load_loss

    def _device_load_balancing(self, scores, top_indices, alpha):
        """
        Device level load balance loss, which regularizes the load from local batch data on each
        expert group to be balanced.
        Please refer to DeepSeek-V2:
        A Strong, Economical, and Efficient Mixture-of-Experts Language Model, https://arxiv.org/abs/2405.04434
        """
        scores = self.reshape(self.transpose_3d(scores, (0, 2, 1)),
                              (scores.shape[0], self.egroup_size, -1))  # (dp, egs, E/egs*N) <- (dp, N, E)
        pi = self.reduce_mean(
            scores, -1)  # (dp, egs) <- (dp, egs, E/egs*N), 1/(E/egs*N) * \sum_{j in E'} \sum_t^N s_{j, t}

        top_indices = self.reshape(self.transpose_3d(top_indices, (0, 2, 1)),
                                   (top_indices.shape[0], -1))  # (dp, kN) <- (dp, N, k)
        mask = self.onehot_2d(top_indices,
                              self.expert_dim,
                              self.on_value,
                              self.off_value)  # (dp, kN, E)fp32 <- (dp, kN)int32
        mask = self.reshape(
            self.transpose_3d(mask, (0, 2, 1)),
            (mask.shape[0], self.egroup_size, -1))  # (dp, egs, E/egs*kN) <- (dp, kN, E)
        # (dp, egs) <- (dp, egs, E/egs*kN), 1/(E/egs*kN) * \sum_{j in E'} \sum_t^N 1(token t selects expert j)
        fi = self.reduce_mean(mask, -1)

        device_load_loss = self.mul(
            self.reduce_mean_2d(
                self.mul_2d(pi, fi)), alpha * self.expert_dim ** 2)  # \alpha * (E/ep)*E \sum_i^{ep} (fi * Pi)

        return device_load_loss

    def _comm_load_balancing(self, scores, top_indices, alpha):
        """
        Communication level load balance loss, which regularizes the load from local batch data on each
        device to be balanced.
        Please refer to DeepSeek-V2:
        A Strong, Economical, and Efficient Mixture-of-Experts Language Model, https://arxiv.org/abs/2405.04434
        """
        scores = self.reshape(
            self.transpose_3d(scores, (0, 2, 1)), (scores.shape[0], self.ep, -1))  # (dp, ep, E/ep*N) <- (dp, N, E)
        pi = self.reduce_mean(scores, -1)  # (dp, ep) <- (dp, ep, E/ep*N), 1/(E/ep*N) * \sum_{j in E'} \sum_t^N s_{j, t}

        top_indices = self.reshape(
            self.transpose_3d(top_indices, (0, 2, 1)), (top_indices.shape[0], -1))  # (dp, kN) <- (dp, N, k)
        mask = self.onehot_2d(top_indices,
                              self.expert_dim,
                              self.on_value,
                              self.off_value)  # (dp, kN, E)fp32 <- (dp, Nk)int32
        mask = self.reshape(
            self.transpose_3d(mask, (0, 2, 1)), (mask.shape[0], self.ep, -1))  # (dp, ep, E/ep*kN) <- (dp, kN, E)
        # (dp, ep) <- (dp, ep, E/dp*kN) ,  1/(E/ep) * \sum_{j in E'} \sum_{t}^N 1(token t sent to device j)
        fi = self.reduce_mean(mask, -1)
        # \alpha * (E/ep)**2 * ep / dp / N \sum_i^{ep} (fi * Pi)
        comm_load_loss = self.mul(self.reduce_mean_2d(self.mul_2d(pi, fi)), alpha * self.expert_dim ** 2)

        return comm_load_loss


class MoEInfer(Cell):
    r"""
        MoEInfer. Routing each tokens to the topk expert and calculating the final output.

        Args:
            ffn (Cell): The FeedForward Module.
            dim (int): The hidden size of each token.
            moe_config (MoEConfig): The configuration of MoE (Mixture of Expert).
            parallel_config: The parallel-related configuration.
        Inputs:
            - **input_tensor** (Tensor) - should be `[batch, seq_length, hidden_size].

        Outputs:
            - **output_tensor** (Tensor) - should be `[batch, seq_length, hidden_size].
    """

    def __init__(self,
                 ffn,
                 dim,
                 moe_config,
                 parallel_config):
        super(MoEInfer, self).__init__()
        self.hidden_size = dim
        self.expert_dim = moe_config.expert_num
        self.topk_norm_prob = moe_config.norm_topk_prob
        self.num_experts_chosen = moe_config.num_experts_chosen
        self.moe_config = moe_config

        self.ffn = ffn
        self.router = Router(d_model=self.hidden_size, moe_config=moe_config, routing_policy=None,
                             training=True, parallel_config=parallel_config)
        self.gating = self.router.dense
        self.noise = 1e-9

        self.zeros = ops.Zeros()
        self.reshape = P.Reshape()
        self.shape = P.Shape()
        self.cast = P.Cast()
        self.mod = P.Mod().shard(((1,), ()))
        self.topk = P.TopK().shard(((1, 1),))
        self.softmax = P.Softmax().shard(((1, 1),))
        self.expand_dims = P.ExpandDims().shard(((1,),))
        self.transpose_2d = P.Transpose().shard(((1, 1),))
        self.sort = P.Sort().shard(((1,),))
        self.gather = P.Gather().shard(((1, 1), (1,)))
        self.onehot = P.OneHot().shard(((1, 1), (), ()))
        self.cumsum = P.CumSum(exclusive=False).shard(((1,),))

        self.on_value = Tensor(1.0, dtype=mstype.float32)
        self.off_value = Tensor(0.0, dtype=mstype.float32)
        self.moe_finalize_routing = MoeFinalizeRouting().shard(((1, 1), (1, 1), (1, 1), (1, 1), (1,), (1, 1)))

    def tensor_sort(self, input_tensor, expert_ids):
        """dispatch and get unsort map for routing"""
        expert_shape = expert_ids.shape
        transposed_index = self.transpose_2d(expert_ids, (1, 0))  # (N, k) -> (k, N)
        reshaped_index = self.reshape(transposed_index, (-1,))  # (k, N) -> (kN)
        _, sort_map = self.sort(reshaped_index)

        inter_map = self.mod(sort_map, expert_shape[0])
        output_tensor = self.gather(input_tensor, inter_map, 0)
        expert_mask = self.onehot(reshaped_index, self.expert_dim, self.on_value, self.off_value)
        expert_cnt = ops.sum(expert_mask, 0)
        group_list = self.cast(self.cumsum(expert_cnt, 0), mstype.int64)

        _, unsort_map = self.sort(sort_map)
        unsort_map = self.cast(unsort_map, mstype.int32)
        return output_tensor, group_list, unsort_map

    def tensor_moe_finalize_routing(self, input_tensor, expert_weight, expert_index, unsort_map):
        """calculate the final output by multiplying FeedForward's output and experts' weight in MoeFinalizeRouting"""
        input_shape = input_tensor.shape  # (2N, h)
        x1 = self.zeros((input_shape[0] // self.num_experts_chosen, input_shape[-1]), input_tensor.dtype)
        x2 = None
        bias = self.zeros((self.expert_dim, input_shape[-1]), input_tensor.dtype)
        expert_weight = self.cast(expert_weight, input_tensor.dtype)
        output_tensor = self.moe_finalize_routing(input_tensor, x1, x2, bias, expert_weight, unsort_map, expert_index)
        return output_tensor

    def construct(self, input_tensor):
        """forward process"""
        input_tensor_shape = self.shape(input_tensor)  # (B, S, H)
        input_dtype = input_tensor.dtype
        input_tensor = self.reshape(
            input_tensor, (-1, self.hidden_size))  # (bs, seq/1, h) -> (bs*seq, h) : use N replace bs*seq

        # gating + topk + softmax
        gating_logits = self.gating(input_tensor.astype(mstype.float32))  # (N, h) * (h, E) -> (bs*seq, E)
        routing_weights = self.softmax(gating_logits.astype(mstype.float32))  # (N, E) -> (N, E)
        expert_val, expert_index = self.topk(routing_weights, self.num_experts_chosen)  # (N, E) -> (N, 2), (N, 2)

        if self.moe_config.norm_topk_prob and self.num_experts_chosen > 1:
            expert_val = self.cast(expert_val, mstype.float32)
            expert_weight = expert_val / (self.expand_dims(ops.sum(expert_val, -1), -1) + 1e-9)
        else:
            expert_weight = ops.mul(self.moe_config.routed_scaling_factor, expert_val)

        expert_weight = self.cast(expert_weight, input_dtype)

        sorted_input_tensor, group_list, unsort_map = self.tensor_sort(input_tensor, expert_index)

        # moeffn
        expert_output = self.ffn(sorted_input_tensor, group_list)  # (N, h) (N, 2) -> (N, 2, h)

        moe_output = self.tensor_moe_finalize_routing(
            expert_output, expert_weight, expert_index, unsort_map)  # -> (N, h)

        output_tensor = self.reshape(moe_output, input_tensor_shape)  # (N, h) -> (bs, seq, h)
        return output_tensor
