Configure attention implementation for FunAudioChat model. Args: config: The model configuration model_args: Model arguments containing flash_attn setting
(config: "PretrainedConfig", model_args)
| 32 | |
| 33 | |
| 34 | def configure_attn_implementation(config: "PretrainedConfig", model_args) -> None: |
| 35 | """ |
| 36 | Configure attention implementation for FunAudioChat model. |
| 37 | |
| 38 | Args: |
| 39 | config: The model configuration |
| 40 | model_args: Model arguments containing flash_attn setting |
| 41 | """ |
| 42 | # Get the flash attention setting from model_args |
| 43 | flash_attn = getattr(model_args, "flash_attn", "auto") |
| 44 | |
| 45 | # Convert string to appropriate value if needed |
| 46 | if isinstance(flash_attn, str): |
| 47 | flash_attn = flash_attn.lower() |
| 48 | |
| 49 | # Handle auto mode - default to FA2 if available |
| 50 | if flash_attn == "auto": |
| 51 | if is_flash_attn_2_available(): |
| 52 | requested_attn_implementation = "flash_attention_2" |
| 53 | logger.info("Auto-selected FlashAttention-2 for FunAudioChat") |
| 54 | else: |
| 55 | logger.warning("FlashAttention-2 not available, using eager attention") |
| 56 | return |
| 57 | |
| 58 | # Handle disabled mode |
| 59 | elif flash_attn == "disabled" or flash_attn == "eager": |
| 60 | requested_attn_implementation = "eager" |
| 61 | |
| 62 | # Handle SDPA mode |
| 63 | elif flash_attn == "sdpa": |
| 64 | if not is_torch_sdpa_available(): |
| 65 | logger.warning("torch>=2.1.1 is required for SDPA attention.") |
| 66 | return |
| 67 | requested_attn_implementation = "sdpa" |
| 68 | |
| 69 | # Handle FA2 mode (primary support) |
| 70 | elif flash_attn == "fa2" or flash_attn == "flash_attention_2": |
| 71 | if not is_flash_attn_2_available(): |
| 72 | logger.warning("FlashAttention-2 is not installed.") |
| 73 | return |
| 74 | requested_attn_implementation = "flash_attention_2" |
| 75 | |
| 76 | else: |
| 77 | logger.warning(f"Unknown attention type: {flash_attn}, using default") |
| 78 | return |
| 79 | |
| 80 | # Configure FunAudioChat model attention |
| 81 | if getattr(config, "model_type", None) == "funaudiochat": |
| 82 | # Audio encoder - only supports FA2 |
| 83 | if hasattr(config, "audio_config"): |
| 84 | setattr(config.audio_config, "_attn_implementation", requested_attn_implementation) |
| 85 | |
| 86 | # Configure CRQ transformer attention if it exists |
| 87 | if hasattr(config.audio_config, "crq_transformer_config") and config.audio_config.crq_transformer_config is not None: |
| 88 | # setattr(config.audo_config, "crq_transformer_attn_implementation", requested_attn_implementation) |
| 89 | config.audio_config.crq_transformer_config["_attn_implementation"] = requested_attn_implementation |
| 90 | |
| 91 | # Text config (LLM backbone) |
nothing calls this directly
no outgoing calls
no test coverage detected