def get_act_ckpt_module(mod_name: str) -> Any:
    """Get the module type from the module name."""
    if mod_name.lower() == 'mptblock':
        mod_type = MPTBlock
    elif mod_name in attention_classes:
        mod_type = attention_classes.get(mod_name)
    elif mod_name.lower() == 'norm_attn_norm':
        mod_type = FusedNormAttentionNorm
    elif mod_name in ffns:
        mod_type = ffns.get(mod_name)
    elif mod_name in ffns_with_norm:
        mod_type = ffns_with_norm.get(mod_name)
    elif mod_name in ffns_with_megablocks:
        mod_type = ffns_with_megablocks.get(mod_name)
    elif mod_name in norms:
        mod_type = norms.get(mod_name)
    else:
        msg = ', '.join(
            list(attention_classes.get_all()) + list(ffns.get_all()) +
            list(ffns_with_norm.get_all()) +
            list(ffns_with_megablocks.get_all()) + list(norms.get_all()) +
            ['MPTBlock'],
        )
        raise ValueError(
            f'{mod_name} (specified in activation_checkpointing_target) is not a recognized option out of available options {msg}.',
        )
    return mod_type
Get the module type from the module name.
llmfoundry/models/utils/act_ckpt.py