def megablocks_n_active_params(mpt_model) -> int:  # type: ignore
    """Calculates the number of active parameters in a MegaBlocks enabled MPT.

    This requires we calculate the number of elements per expert and
    multiply this by top k.

    Args:
        mpt_model (ComposerMPTCausalLM): MPT model of which the number of
            active parameters is calculated.

    Returns:
        An int for the active number of parameters in this MPT model.
    """
    import megablocks

    moe_num_experts = mpt_model.config.ffn_config.get('moe_num_experts', 1)
    moe_world_size = mpt_model.config.ffn_config.get('moe_world_size')

    local_experts = moe_num_experts / moe_world_size  # if local_experts is < 1, then the expert is sharded

    moe_top_k = mpt_model.config.ffn_config.get('moe_top_k', 1)
    n_active_params = 0
    for module in mpt_model.modules():
        if isinstance(
            module,
            (megablocks.layers.mlp.SparseMLP, megablocks.layers.mlp.MLP),
        ):
            n_w1 = _dtensor_safe_check_numel(module.w1)
            n_active_params += int(n_w1 / local_experts * moe_top_k)
            n_w2 = _dtensor_safe_check_numel(module.w2)
            n_active_params += int(n_w2 / local_experts * moe_top_k)

            # GLU has an extra weight
            if hasattr(module, 'v1'):
                n_v1 = _dtensor_safe_check_numel(module.v1)
                n_active_params += int(n_v1 / local_experts * moe_top_k)
        else:
            n_active_params += module_n_params(module)

    return n_active_params
Calculates the number of active parameters in a MegaBlocks enabled MPT.

    This requires we calculate the number of elements per expert and
    multiply this by top k.

    Args:
        mpt_model (ComposerMPTCausalLM): MPT model of which the number of
            active parameters is calculated.

    Returns:
        An int for the active number of parameters in this MPT model.
    
llmfoundry/models/utils/mpt_param_count.py