Register FunAudioChat attention configuration plugin with LLaMA-Factory. This plugin handles FlashAttention-2 configuration for FunAudioChat models.
()
| 49 | |
| 50 | |
| 51 | def register_attention_plugin() -> None: |
| 52 | """ |
| 53 | Register FunAudioChat attention configuration plugin with LLaMA-Factory. |
| 54 | This plugin handles FlashAttention-2 configuration for FunAudioChat models. |
| 55 | """ |
| 56 | try: |
| 57 | # Import the attention module |
| 58 | from . import attention |
| 59 | |
| 60 | # Try to inject into llamafactory's model_utils |
| 61 | try: |
| 62 | from llamafactory.model.model_utils import attention as llama_attention |
| 63 | from llamafactory.model import patcher |
| 64 | |
| 65 | # Get the original configure function |
| 66 | original_configure = llama_attention.configure_attn_implementation |
| 67 | original_print = llama_attention.print_attn_implementation |
| 68 | |
| 69 | # Create wrapper functions that handle both original and FunAudioChat models |
| 70 | def configure_attn_implementation_wrapper(config, model_args): |
| 71 | # First try FunAudioChat configuration |
| 72 | if getattr(config, "model_type", None) == "funaudiochat": |
| 73 | attention.configure_attn_implementation(config, model_args) |
| 74 | else: |
| 75 | # Fall back to original implementation |
| 76 | original_configure(config, model_args) |
| 77 | |
| 78 | def print_attn_implementation_wrapper(config): |
| 79 | # First try FunAudioChat print |
| 80 | if getattr(config, "model_type", None) == "funaudiochat": |
| 81 | attention.print_attn_implementation(config) |
| 82 | else: |
| 83 | # Fall back to original implementation |
| 84 | original_print(config) |
| 85 | |
| 86 | # Replace the functions in both llamafactory.model.model_utils.attention |
| 87 | # and llamafactory.model.patcher (which imports them directly) |
| 88 | llama_attention.configure_attn_implementation = configure_attn_implementation_wrapper |
| 89 | llama_attention.print_attn_implementation = print_attn_implementation_wrapper |
| 90 | patcher.configure_attn_implementation = configure_attn_implementation_wrapper |
| 91 | patcher.print_attn_implementation = print_attn_implementation_wrapper |
| 92 | |
| 93 | logger.info("Registered attention plugin for FunAudioChat") |
| 94 | |
| 95 | except ImportError as e: |
| 96 | logger.warning(f"Could not inject attention plugin into llamafactory: {e}") |
| 97 | logger.warning("FunAudioChat attention configuration may not work properly") |
| 98 | |
| 99 | except Exception as e: |
| 100 | logger.error(f"Failed to register attention plugin: {e}") |
| 101 | raise |
| 102 | |
| 103 | |
| 104 | def register_templates() -> None: |