def mpt_get_total_params(mpt_model) -> int:  # type: ignore
    """Calculates the total parameter count of an MPT model.

    Note: Must be called before model parameters are sharded by FSDP.

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

    Returns:
        An int for the total number of parameters in this MPT model.
    """
    if mpt_model.config.ffn_config['ffn_type'] in ffns_with_megablocks:
        return megablocks_n_total_params(mpt_model)
    else:
        return sum(p.numel() for p in mpt_model.parameters())
Calculates the total parameter count of an MPT model.

    Note: Must be called before model parameters are sharded by FSDP.

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

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