Function print_attn_implementation

training/plugin/attention.py:105–130

Print the attention implementation being used. Args: config: The model configuration

(config: "PretrainedConfig")

def print_attn_implementation(config: "PretrainedConfig") -> None:
    """
    Print the attention implementation being used.
    
    Args:
        config: The model configuration
    """
    attn_implementation = getattr(config, "_attn_implementation", None)
    
    if attn_implementation == "flash_attention_2":
        logger.info("Using FlashAttention-2 for faster training and inference.")
    elif attn_implementation == "sdpa":
        logger.info("Using torch SDPA for faster training and inference.")
    else:
        logger.info("Using vanilla attention implementation.")
    
    # Print sub-config attention if it's a FunAudioChat model
    if getattr(config, "model_type", None) == "funaudiochat":
        if hasattr(config, "audio_config"):
            audio_attn = getattr(config.audio_config, "_attn_implementation", None)
            if audio_attn:
                logger.info(f"  - Audio encoder attention: {audio_attn}")
        if hasattr(config, "text_config"):
            text_attn = getattr(config.text_config, "_attn_implementation", None)
            if text_attn:
                logger.info(f"  - Text config attention: {text_attn}")


__all__ = ["configure_attn_implementation", "print_attn_implementation"]