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

    MoE experts are sharded across workers. This function scans for MegaBlocks
    modules then multiplies expert params count by MoE world size.

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

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

    moe_world_size = mpt_model.config.ffn_config.get('moe_world_size')

    n_total_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_total_params += n_w1 * moe_world_size
            n_w2 = _dtensor_safe_check_numel(module.w2)
            n_total_params += n_w2 * moe_world_size

            # GLU has an extra weight
            if hasattr(module, 'v1'):
                n_v1 = _dtensor_safe_check_numel(module.v1)
                n_total_params += n_v1 * moe_world_size
        else:
            n_total_params += module_n_params(module)

    return n_total_params
Calculates the number of parameters in a MegaBlocks enabled MPT model.

    MoE experts are sharded across workers. This function scans for MegaBlocks
    modules then multiplies expert params count by MoE world size.

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

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