def mpt_get_active_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 active number of parameters in this MPT model.
    """
    if mpt_model.config.ffn_config['ffn_type'] in ffns_with_megablocks:
        params = megablocks_n_active_params(mpt_model)
    else:
        params = sum(p.numel() for p in mpt_model.parameters())
    if not mpt_model.model.transformer.config.tie_word_embeddings:
        # Embedding layers are lookup tables, therefore are not counted in the FLOP computation
        params -= _dtensor_safe_check_numel(
            mpt_model.model.transformer.wte.weight,
        )
    return params
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 active number of parameters in this MPT model.
    
llmfoundry/models/utils/mpt_param_count.py