def estimate_moe_param(
    vocab_size,
    hidden_size,
    num_hidden_layers,
    intermediate_size,
    num_experts,
    num_selects,
):
    """
    Llama Structure with SwiGLU MLP
    MoE split on intermediate_size without weight noise in 2-layer gate
    """

    emb = vocab_size * hidden_size
    lm_head = vocab_size * hidden_size
    final_norm = hidden_size

    self_attn = hidden_size * hidden_size * 4
    mlp = hidden_size * intermediate_size * 3
    input_norm = hidden_size
    post_attn_norm = hidden_size

    dense_one_layer = self_attn + mlp + input_norm + post_attn_norm
    dense_mid = dense_one_layer * num_hidden_layers
    dense_params = emb + lm_head + final_norm + dense_mid

    gate = hidden_size * num_experts + num_experts * num_selects
    moe_one_layer = self_attn + mlp + input_norm + post_attn_norm + gate
    moe_mid = moe_one_layer * num_hidden_layers
    moe_total_params = emb + lm_head + final_norm + moe_mid

    moe_one_act_layer = (
        self_attn
        + (mlp / num_experts * num_selects)
        + input_norm
        + post_attn_norm
        + gate
    )
    moe_act_mid = moe_one_act_layer * num_hidden_layers
    moe_act_params = emb + lm_head + final_norm + moe_act_mid

    return {
        "dense_params": dense_params,
        "moe_total_params": moe_total_params,
        "moe_act_params": moe_act_params,
        "dense_mid": dense_mid,
        "moe_mid": moe_mid,
        "moe_act_mid": moe_act_mid,
    }

    Llama Structure with SwiGLU MLP
    MoE split on intermediate_size without weight noise in 2-layer gate
    
smoe/utils/param_estimation.py