Class SparseMLP

analysis/colossalai_replace/layer.py:20–402

A class for users to create MoE modules in their models. Args: dim_model (int): Hidden dimension of training model num_experts (int): The number experts top_k (int, optional): The number of experts for dispatchment of each token capacity_factor_train (float, opti

import numpy as np

class SparseMLP(nn.Module):
    """A class for users to create MoE modules in their models.

    Args:
        dim_model (int): Hidden dimension of training model
        num_experts (int): The number experts
        top_k (int, optional): The number of experts for dispatchment of each token
        capacity_factor_train (float, optional): Capacity factor in routing during training
        capacity_factor_eval (float, optional): Capacity factor in routing during evaluation
        min_capacity (int, optional): The minimum number of the capacity of each expert
        noisy_policy (str, optional): The policy of noisy function. Now we have 'Jitter' and 'Gaussian'.
            'Jitter' can be found in `Switch Transformer paper`_.
            'Gaussian' can be found in `ViT-MoE paper`_.
        drop_tks (bool, optional): Whether drops tokens in evaluation
        use_residual (bool, optional): Makes this MoE layer a Residual MoE.
            More information can be found in `Microsoft paper`_.
        residual_instance (nn.Module, optional): The instance of residual module in Residual MoE
        expert_instance (MoeExperts, optional): The instance of experts module in MoeLayer
        expert_cls (Type[nn.Module], optional): The class of each expert when no instance is given
        expert_args (optional): The args of expert when no instance is given

    .. _Switch Transformer paper:
        https://arxiv.org/abs/2101.03961
    .. _ViT-MoE paper:
        https://arxiv.org/abs/2106.05974
    .. _Microsoft paper:
        https://arxiv.org/abs/2201.05596
    """

    def __init__(
        self,
        num_experts: int,
        hidden_size: int,
        intermediate_size: int,
        router_top_k: int = 1,
        router_capacity_factor_train: float = 1.25,
        router_capacity_factor_eval: float = 2.0,
        router_min_capacity: int = 4,
        router_noisy_policy: Optional[str] = None,
        router_drop_tks: bool = True,
        mlp_activation: Optional[str] = None,
        mlp_gated: bool = False,
        enable_load_balance: bool = False,
        load_balance_tolerance: float = 0.1,
        load_balance_beam_width: int = 8,
        load_balance_group_swap_factor: float = 0.4,
        enable_kernel: bool = False,
        enable_comm_overlap: bool = False,
        enable_hierarchical_comm: bool = False,
        model_output_dir: str = None,
    ):
        super().__init__()
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size
        self.num_experts = num_experts
        self.gated = mlp_gated
        self.enable_kernel = enable_kernel
        self.enable_comm_overlap = enable_comm_overlap