from collections import namedtuple
from typing import Any, Dict, Iterable, Union, List, Optional, Tuple
import math
import torch
from torch import nn
from vllm import _mlu_ops as mlu_ops
from vllm.config import CacheConfig, VllmConfig
from vllm_mlu.transformers_utils.configs import CustomConfig
from vllm.attention import Attention, AttentionMetadata
from vllm.distributed import (get_pp_group, get_tensor_model_parallel_rank,
                              get_tensor_model_parallel_world_size,
                              tensor_model_parallel_all_reduce)
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.linear import (QKVParallelLinear,
                                               RowParallelLinear,
                                               ReplicatedLinear)
from vllm.model_executor.layers.logits_processor import LogitsProcessor
from vllm.model_executor.layers.quantization.base_config import (
    QuantizationConfig)
from vllm.model_executor.layers.rotary_embedding import get_rope
from vllm.model_executor.layers.sampler import Sampler, SamplerOutput
from vllm.model_executor.layers.vocab_parallel_embedding import (
    DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)
from vllm.model_executor.models.interfaces import SupportsPP
from vllm.model_executor.sampling_metadata import SamplingMetadata
from vllm.sequence import IntermediateTensors
from vllm_mlu.model_executor.layers.feed_forward import FeedForward
from vllm.model_executor.models.utils import PPMissingLayer, make_layers
from vllm_mlu.model_executor.models.layer_utils import (
    decoder_layer_forward_base, is_per_tensor_smoothquant,
    is_per_token_smoothquant, quant_fusion_with_rmsnorm,
    quant_fusion_with_layernorm)


def _get_alibi_slopes(total_num_heads: int) -> torch.Tensor:
    closest_power_of_2 = 2**math.floor(math.log2(total_num_heads))
    base = torch.tensor(
        2**(-(2**-(math.log2(closest_power_of_2) - 3))),
        dtype=torch.float32,
    )
    powers = torch.arange(1, 1 + closest_power_of_2, dtype=torch.int32)
    slopes = torch.pow(base, powers)

    if closest_power_of_2 != total_num_heads:
        extra_base = torch.tensor(
            2**(-(2**-(math.log2(2 * closest_power_of_2) - 3))),
            dtype=torch.float32,
        )
        num_remaining_heads = min(closest_power_of_2,
                                  total_num_heads - closest_power_of_2)
        extra_powers = torch.arange(start=1,
                                    end=1 + 2 * num_remaining_heads,
                                    step=2,
                                    dtype=torch.int32)
        slopes = torch.cat(
            [slopes, torch.pow(extra_base, extra_powers)], dim=0)
    return slopes


class LayerNorm(nn.Module):

    def __init__(
        self,
        hidden_size: int,
        eps: float = 1e-6,
    ) -> None:
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.bias = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps

    def forward(
        self,
        x: torch.Tensor,
        residual: Optional[torch.Tensor] = None,
    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
        x = x.view(-1, self.weight.data.shape[0])
        if residual is not None:
            residual = residual.view(-1, self.weight.data.shape[0])
            return mlu_ops.fused_layer_norm(x, residual, self.weight.data, self.bias.data, None, self.variance_epsilon, True)
        else:
            return mlu_ops.fused_layer_norm(x, residual, self.weight.data, self.bias.data, None, self.variance_epsilon, False)


_NORM_DICT: Dict[str, nn.Module] = {"rmsnorm": RMSNorm, "layernorm": LayerNorm}


class CustomMoeBlock(nn.Module):

    def __init__(
        self,
        config: CustomConfig,
        quant_config: Optional[QuantizationConfig] = None,
    ):
        super().__init__()
        self.config = config
        self.rank = get_tensor_model_parallel_rank()
        self.tp_size = get_tensor_model_parallel_world_size()
        self.n_routed_experts = config.num_experts
        self.top_k = config.num_experts_per_tok
        if self.tp_size > self.n_routed_experts:
            raise ValueError(
                f"Tensor parallel size {self.tp_size} is greater than "
                f"the number of experts {self.n_routed_experts}.")

        self.moe_intermediate_size = self.config.moe_intermediate_size // self.tp_size

        if quant_config is None:
            self.w1 = nn.Parameter(
                torch.empty(self.config.num_experts,
                            2 * self.moe_intermediate_size if self.config.is_gated else self.moe_intermediate_size,
                            self.config.hidden_size,
                            dtype=torch.get_default_dtype()), requires_grad=False)
            self.w2 = nn.Parameter(
                torch.empty(self.config.num_experts,
                            self.config.hidden_size,
                            self.moe_intermediate_size,
                            dtype=torch.get_default_dtype()), requires_grad=False)
            self.w1_scale = None
            self.w2_scale = None
            self.input_smooth = None
            self.act_smooth = None
        else:
            assert quant_config.weight_bits == 8
            self.w1 = nn.Parameter(
                torch.empty(self.config.num_experts,
                            2 * self.moe_intermediate_size if self.config.is_gated else self.moe_intermediate_size,
                            self.config.hidden_size,
                            device="mlu",
                            dtype=torch.int8), requires_grad=False)
            self.w2 = nn.Parameter(
                torch.empty(self.config.num_experts,
                            self.config.hidden_size,
                            self.moe_intermediate_size,
                            device="mlu",
                            dtype=torch.int8), requires_grad=False)
            self.w1_scale = nn.Parameter(
                torch.empty(
                    self.config.num_experts,
                    2 * self.moe_intermediate_size if self.config.is_gated else self.moe_intermediate_size,
                    device="mlu",
                    dtype=torch.float32), requires_grad=False)
            self.w2_scale = nn.Parameter(
                torch.empty(
                    self.config.num_experts,
                    self.config.hidden_size,
                    device="mlu",
                    dtype=torch.float32), requires_grad=False)
            self.input_smooth = None
            self.act_smooth = None
            if quant_config.quant_mode == "SmoothQuant":
                self.input_smooth =nn.Parameter(
                    torch.empty(
                        self.config.num_experts,
                        self.config.hidden_size,
                        device="mlu",
                        dtype=torch.float32), requires_grad=False)
                self.act_smooth =nn.Parameter(
                    torch.empty(
                        self.config.num_experts,
                        self.moe_intermediate_size,
                        device="mlu",
                        dtype=torch.float32), requires_grad=False)

        self.gate = ReplicatedLinear(config.hidden_size,
                                     self.n_routed_experts,
                                     bias=False,
                                     quant_config=None)
        if config.shared_expert_intermediate_size > 0:
            self.shared_expert = FeedForward(hidden_size=config.hidden_size,
                                             intermediate_size=config.shared_expert_intermediate_size,
                                             hidden_act = self.config.hidden_act,
                                             up_proj_name='gate_up_proj',
                                             is_gated=self.config.is_gated,
                                             down_proj_name='down_proj',
                                             bias=False,
                                             quant_config=quant_config,
                                             reduce_results=False)
        else:
            self.shared_expert = None
        self.shared_expert_gate = torch.nn.Linear(config.hidden_size,
                                                  1,
                                                  bias=False)

    def forward(self, hidden_states: torch.Tensor, residual: Optional[torch.Tensor] = None) -> torch.Tensor:
        num_tokens, hidden_dim = hidden_states.shape
        hidden_states = hidden_states.view(-1, hidden_dim)
        shared_output = None
        if self.shared_expert is not None:
            shared_output = self.shared_expert(hidden_states)
            if self.shared_expert_gate is not None:
                shared_output = F.sigmoid(
                    self.shared_expert_gate(hidden_states)) * shared_output

        # router_logits: (num_tokens, n_experts)
        router_logits, _ = self.gate(hidden_states)
        residual_ = None if self.rank > 0 else residual
        final_hidden_states = mlu_ops.fused_moe(hidden_states,
                                               router_logits,
                                               self.w1,
                                               self.w2,
                                               None,
                                               None,
                                               residual_,
                                               self.input_smooth,
                                               self.act_smooth,
                                               self.w1_scale,
                                               self.w2_scale,
                                               self.top_k,
                                               self.config.norm_topk_prob,
                                               self.config.is_gated,
                                               self.config.hidden_act)
        if shared_output is not None:
            final_hidden_states = final_hidden_states + shared_output

        reduce_results = (self.config.use_parallel_residual == False)
        if reduce_results:
            final_hidden_states = tensor_model_parallel_all_reduce(
                final_hidden_states)

        return final_hidden_states.view(num_tokens, hidden_dim)


class CustomAttention(nn.Module):

    def __init__(
        self,
        config: CustomConfig,
        cache_config: Optional[CacheConfig] = None,
        quant_config: Optional[QuantizationConfig] = None,
    ) -> None:
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size
        attention_bias = getattr(config, "attention_bias", False) or getattr(config, "bias", False)
        self.hidden_size = config.hidden_size
        tp_size = get_tensor_model_parallel_world_size()
        self.total_num_heads = config.num_attention_heads
        assert self.total_num_heads % tp_size == 0
        self.num_heads = self.total_num_heads // tp_size
        num_kv_heads=getattr(config, "num_key_value_heads", config.num_attention_heads)
        self.total_num_kv_heads = num_kv_heads
        if self.total_num_kv_heads >= tp_size:
            # Number of KV heads is greater than TP size, so we partition
            # the KV heads across multiple tensor parallel GPUs.
            assert self.total_num_kv_heads % tp_size == 0
        else:
            # Number of KV heads is less than TP size, so we replicate
            # the KV heads across multiple tensor parallel GPUs.
            assert tp_size % self.total_num_kv_heads == 0
        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
        self.head_dim = self.hidden_size // self.total_num_heads
        self.q_size = self.num_heads * self.head_dim
        self.kv_size = self.num_kv_heads * self.head_dim
        self.kv_scale = 1.0
        self.scaling = self.head_dim**-0.5

        self.qkv_proj = QKVParallelLinear(
            self.hidden_size,
            self.head_dim,
            self.total_num_heads,
            self.total_num_kv_heads,
            bias=attention_bias,
            quant_config=quant_config,
        )
        self.o_proj = RowParallelLinear(
            self.total_num_heads * self.head_dim,
            self.hidden_size,
            bias=attention_bias,
            quant_config=quant_config,
            skip_bias_add=(self.config.use_parallel_residual and attention_bias),
            reduce_results = (self.config.use_parallel_residual == False),
        )

        self.alibi_slopes = None
        self.rotary_emb = None
        if self.config.position_embedding_type == "ALIBI":
            tp_rank = get_tensor_model_parallel_rank()
            head_start = tp_rank * self.num_heads
            head_end = (tp_rank + 1) * self.num_heads
            alibi_slopes = _get_alibi_slopes(self.total_num_heads)
            self.alibi_slopes = alibi_slopes[head_start:head_end].tolist()
        else:
            rope_theta = getattr(config, "rope_theta", 10000)
            rope_scaling = getattr(config, "rope_scaling", None)
            if rope_scaling is not None and getattr(
                    config, "original_max_position_embeddings", None):
                rope_scaling["original_max_position_embeddings"] = (
                    config.original_max_position_embeddings)
            max_position_embeddings = getattr(config, "max_sequence_length", 8192)
            is_neox_style = getattr(config, "is_neox_style", False)
            self.rotary_emb = get_rope(
                self.head_dim,
                rotary_dim=self.head_dim,
                max_position=max_position_embeddings,
                base=rope_theta,
                is_neox_style=is_neox_style,
                rope_scaling=rope_scaling,
            )

        self.attn = Attention(self.num_heads,
                              self.head_dim,
                              self.scaling,
                              num_kv_heads=self.num_kv_heads,
                              alibi_slopes=self.alibi_slopes,
                              cache_config=cache_config)

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
        residual: Optional[torch.Tensor] = None,
        smooth_quant_scale: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        qkv, _ = self.qkv_proj(hidden_states, smooth_quant_scale)
        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
        if self.rotary_emb:
            qk, _ = qkv.split([self.q_size + self.kv_size, self.kv_size], dim=-1)
            self.rotary_emb(positions, qk.view(-1, self.num_heads + self.num_kv_heads, self.head_dim))
        attn_output = self.attn(q, k, v, kv_cache, attn_metadata)
        output, bias = self.o_proj(attn_output, residual)
        if self.o_proj.skip_bias_add and get_tensor_model_parallel_rank() == 0:
            output += bias
        return output


class CustomDecoderLayer(nn.Module):

    def __init__(
        self,
        config: CustomConfig,
        cache_config: Optional[CacheConfig] = None,
        quant_config: Optional[QuantizationConfig] = None,
    ) -> None:
        super().__init__()
        self.config = config
        self.self_attn = CustomAttention(
            config=config,
            cache_config=cache_config,
            quant_config=quant_config,
        )

        mlp_bias = getattr(config, "mlp_bias", False) or getattr(config, "bias", False)
        is_gated = getattr(config, "is_gated", False)

        if config.num_experts is not None:
            self.mlp = CustomMoeBlock(config=config,
                                    quant_config=quant_config)
        else:
            self.mlp = FeedForward(hidden_size=config.hidden_size,
                                   intermediate_size=config.intermediate_size,
                                   hidden_act=self.config.hidden_act,
                                   up_proj_name='up_proj',
                                   is_gated=is_gated,
                                   down_proj_name='down_proj',
                                   bias=mlp_bias,
                                   quant_config=quant_config,
                                   skip_bias_add=(self.config.use_parallel_residual and mlp_bias),
                                   reduce_results = (self.config.use_parallel_residual == False))

        self.input_layernorm = _NORM_DICT[self.config.norm_type](config.hidden_size, eps=config.norm_eps)
        self.post_attention_layernorm = _NORM_DICT[self.config.norm_type](config.hidden_size, eps=config.norm_eps)

        # perf per-tensor sq cases by fusing quantization in layernorm
        self.is_per_tesnor_sq_perf_cases = (is_per_tensor_smoothquant(quant_config) and
                                            not self.config.apply_residual_connection_post_layernorm)
        self.is_per_token_sq_perf_cases = (is_per_token_smoothquant(quant_config) and
                                            not self.config.apply_residual_connection_post_layernorm)
        if self.is_per_tesnor_sq_perf_cases or self.is_per_token_sq_perf_cases:
            self.self_attn.qkv_proj.quant_method.skip_quant_input = True
            self.quant_fusion_attn_layernorm = None
            self.is_moe = config.num_experts is not None
            self.use_rmsnorm = self.config.norm_type == "rmsnorm"
            if not self.is_moe:
                self.mlp.up_proj.quant_method.skip_quant_input = True
                self.quant_fusion_mlp_layernorm = None

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
        residual: Optional[torch.Tensor],
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        if self.config.use_parallel_residual:
            # x = x + attn(ln1(x)) + mlp(ln2(x))
            layernorm_output = self.input_layernorm(hidden_states)
            attention_output = self.self_attn(
                positions=positions,
                hidden_states=layernorm_output,
                kv_cache=kv_cache,
                attn_metadata=attn_metadata,
            )

            layernorm_output = self.post_attention_layernorm(hidden_states)
            if self.mlp.skip_bias_add:
                mlp_output, mlp_bias = self.mlp(layernorm_output)
                if get_tensor_model_parallel_rank() == 0:
                    mlp_output += mlp_bias
            else:
                mlp_output = self.mlp(layernorm_output)

            if get_tensor_model_parallel_rank() == 0:
                hidden_states = mlp_output + attention_output + hidden_states
            else:
                hidden_states = mlp_output + attention_output

            hidden_states = tensor_model_parallel_all_reduce(hidden_states)
            return hidden_states, None
        else:
            # rmsnorm use fused_rms_norm to get better performance
            # if apply_residual_connection_post_layernorm:
            #     x = ln1(x) + attn(ln1(x))
            #     x = ln2(x) + mlp(ln2(x))
            # else:
            #     x = x + attn(ln1(x))
            #     x = x + mlp(ln2(x))
            attn_layernorm = self.input_layernorm
            mlp_layernorm = self.post_attention_layernorm
            if self.is_per_tesnor_sq_perf_cases:
                quant_fusion_func = (quant_fusion_with_rmsnorm if
                                     self.use_rmsnorm else quant_fusion_with_layernorm)
                if self.quant_fusion_attn_layernorm is None:
                    self.quant_fusion_attn_layernorm = quant_fusion_func(
                        self.input_layernorm, self.self_attn.qkv_proj.scale_to_int)
                attn_layernorm = self.quant_fusion_attn_layernorm
                if not self.is_moe:
                    if self.quant_fusion_mlp_layernorm is None:
                        self.quant_fusion_mlp_layernorm = quant_fusion_func(
                            self.post_attention_layernorm, self.mlp.up_proj.scale_to_int)
                    mlp_layernorm = self.quant_fusion_mlp_layernorm
            elif self.is_per_token_sq_perf_cases:
                quant_fusion_func = (quant_fusion_with_rmsnorm if
                                     self.use_rmsnorm else quant_fusion_with_layernorm)
                if self.quant_fusion_attn_layernorm is None:
                    self.quant_fusion_attn_layernorm = quant_fusion_func(
                        self.input_layernorm, self.self_attn.qkv_proj.smooth, dynamic_quant=True)
                attn_layernorm = self.quant_fusion_attn_layernorm
                if not self.is_moe:
                    if self.quant_fusion_mlp_layernorm is None:
                        self.quant_fusion_mlp_layernorm = quant_fusion_func(
                            self.post_attention_layernorm, self.mlp.up_proj.smooth, dynamic_quant=True)
                    mlp_layernorm = self.quant_fusion_mlp_layernorm

            post_norm_fuse_en=(self.is_per_token_sq_perf_cases and not self.is_moe)
            return decoder_layer_forward_base(positions=positions,
                                      hidden_states=hidden_states,
                                      kv_cache=kv_cache,
                                      attn_metadata=attn_metadata,
                                      input_layernorm=attn_layernorm,
                                      self_attn=self.self_attn,
                                      post_layernorm=mlp_layernorm,
                                      mlp=self.mlp,
                                      apply_residual_connection_post_layernorm=self.config.apply_residual_connection_post_layernorm,
                                      input_norm_fuse_en=self.is_per_token_sq_perf_cases,
                                      post_norm_fuse_en=post_norm_fuse_en), None


class CustomModel(nn.Module):

    def __init__(
        self,
        config: CustomConfig,
        cache_config: Optional[CacheConfig] = None,
        quant_config: Optional[QuantizationConfig] = None,
    ) -> None:
        super().__init__()
        self.config = config
        if get_pp_group().is_first_rank or (config.tie_word_embeddings
                                            and get_pp_group().is_last_rank):
            embed_layer = VocabParallelEmbedding if self.config.use_parallel_embedding else nn.Embedding
            self.embed_tokens = embed_layer(
                config.vocab_size,
                config.hidden_size,
            )
        else:
            self.embed_tokens = PPMissingLayer()

        self.start_layer, self.end_layer, self.layers = make_layers(
            config.num_hidden_layers,
            lambda prefix: CustomDecoderLayer(config=config,
                                             cache_config=cache_config,
                                             quant_config=quant_config),
            prefix="custom_model")

        if get_pp_group().is_last_rank:
            self.norm = _NORM_DICT[self.config.norm_type](config.hidden_size, eps=config.norm_eps)
        else:
            self.norm = PPMissingLayer()

    def forward(
        self,
        input_ids: Optional[torch.Tensor],
        positions: torch.Tensor,
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
        intermediate_tensors: Optional[IntermediateTensors] = None,
        inputs_embeds: Optional[torch.Tensor] = None,
    ) -> Union[torch.Tensor, IntermediateTensors]:
        if get_pp_group().is_first_rank:
            if inputs_embeds is not None:
                hidden_states = inputs_embeds
            else:
                hidden_states = self.embed_tokens(input_ids)
            residual = None
        else:
            assert intermediate_tensors is not None
            hidden_states = intermediate_tensors["hidden_states"]
            residual = intermediate_tensors["residual"]

        for i in range(self.start_layer, self.end_layer):
            layer = self.layers[i]
            hidden_states, residual = layer(
                positions,
                hidden_states,
                kv_caches[i - self.start_layer],
                attn_metadata,
                residual,
            )

        if not get_pp_group().is_last_rank:
            return IntermediateTensors({
                "hidden_states": hidden_states,
                "residual": residual
            })

        if residual is not None:
            hidden_states, _ = self.norm(hidden_states, residual)
        else:
            hidden_states = self.norm(hidden_states)
        return hidden_states


class CustomForCausalLM(nn.Module, SupportsPP):

    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
        super().__init__()
        self.config = vllm_config.model_config.hf_text_config
        self.quant_config = vllm_config.quant_config
        self.cache_config = vllm_config.cache_config
        self._verify_params()
        self.model = CustomModel(self.config, self.cache_config, self.quant_config)

        if get_pp_group().is_last_rank:
            self.lm_head = ParallelLMHead(self.config.vocab_size, self.config.hidden_size)
            self.logits_processor = LogitsProcessor(self.config.vocab_size)
            self.sampler = Sampler()
        else:
            self.lm_head = PPMissingLayer()

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
        intermediate_tensors: Optional[IntermediateTensors] = None,
    ) -> Union[torch.Tensor, IntermediateTensors]:
        hidden_states = self.model(input_ids, positions, kv_caches,
                                   attn_metadata, intermediate_tensors)
        return hidden_states

    def compute_logits(self, hidden_states: torch.Tensor,
                       sampling_metadata: SamplingMetadata) -> torch.Tensor:
        logits = self.logits_processor(self.lm_head, hidden_states,
                                       sampling_metadata)
        return logits

    def sample(
        self,
        logits: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[SamplerOutput]:
        next_tokens = self.sampler(logits, sampling_metadata)
        return next_tokens
    
    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
        pass

    def make_empty_intermediate_tensors(
            self, batch_size: int, dtype: torch.dtype,
            device: torch.device) -> IntermediateTensors:
        return IntermediateTensors({
            "hidden_states":
            torch.zeros((batch_size, self.config.hidden_size),
                        dtype=dtype,
                        device=device),
            "residual":
            torch.zeros((batch_size, self.config.hidden_size),
                        dtype=dtype,
                        device=device),
        })

    def _verify_params(self) -> None:
        if (self.config.max_sequence_length) is None or \
           (self.config.num_hidden_layers) is None or \
           (self.config.hidden_size) is None or \
           (self.config.vocab_size) is None or \
           (self.config.num_attention_heads) is None:
            raise ValueError(
                "max_sequence_length, num_hidden_layers, hidden_size, vocab_size, "
                "num_attention_heads, must be vaild int values")

        if self.config.hidden_act not in ["silu", "gelu"]:
            raise ValueError(
                "CustomConfig hidden_act must be one of [silu, gelu]. Got "
                f"{self.config.hidden_act}.")

        if self.config.position_embedding_type not in ["ALIBI", "ROPE"]:
            raise ValueError(
                "position_embedding_type must be one of [ALIBI, ROPE]. Got "
                f"{self.config.position_embedding_type}.")

        if self.config.num_experts is not None:
            if self.config.num_experts_per_tok is None:
                raise ValueError(
                    "num_experts_per_tok must be a valid int value when num_experts is not None")
            if self.config.moe_intermediate_size is None:
                raise ValueError(
                    "moe_intermediate_size must be a valid int value when num_experts is not None")
            if self.config.shared_expert_intermediate_size is None:
                raise ValueError(
                    "shared_expert_intermediate_size must be a valid int value when num_experts is not None")
            if self.config.norm_topk_prob is None:
                raise ValueError(
                    "norm_topk_prob must be a valid bool value when num_experts is not None")
            if self.config.mlp_bias is True:
                raise ValueError(
                    "mlp_bias must be False when num_experts is not None")
            if self.quant_config is not None and self.quant_config.get_name() != "SmoothQuant":
                raise ValueError(
                    "moe only support smoothquant now")
        else:
            if self.config.intermediate_size is None:
                raise ValueError(
                    "intermediate_size must be a valid int value when num_experts is None")

        if self.config.norm_type not in ["rmsnorm", "layernorm"]:
            raise ValueError(
                "norm_type must be one of [rmsnorm, layernorm]. Got "
                f"{self.config.norm_type}.")
