mirror of https://github.com/razor-ai/soup.git
130 lines
3.5 KiB
Python
130 lines
3.5 KiB
Python
"""FlashAttention auto-detection and configuration.
|
|
|
|
Detects FlashAttention availability (v2/v3/v4) and configures models
|
|
to use the best available attention implementation automatically.
|
|
|
|
FlashAttention provides 2-4x speedup and significant memory savings
|
|
for long sequences by avoiding materializing the full attention matrix.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
# Ordered by preference (newest first)
|
|
FLASH_ATTN_VERSIONS = ("flash_attention_3", "flash_attention_2")
|
|
|
|
|
|
def check_flash_attn_available() -> str | None:
|
|
"""Detect the best available FlashAttention implementation.
|
|
|
|
Returns:
|
|
The attention implementation string for model_kwargs, or None if unavailable.
|
|
One of: "flash_attention_3", "flash_attention_2", None.
|
|
"""
|
|
# FlashAttention requires CUDA
|
|
try:
|
|
import torch
|
|
|
|
if not torch.cuda.is_available():
|
|
return None
|
|
except ImportError:
|
|
return None
|
|
|
|
# Check FlashAttention 3 (Hopper architecture, H100+)
|
|
try:
|
|
import flash_attn # noqa: F401
|
|
|
|
version = getattr(flash_attn, "__version__", "0.0.0")
|
|
major = int(version.split(".")[0])
|
|
if major >= 3:
|
|
return "flash_attention_3"
|
|
except (ImportError, ValueError, IndexError):
|
|
pass
|
|
|
|
# Check FlashAttention 2
|
|
try:
|
|
from transformers.utils import is_flash_attn_2_available
|
|
|
|
if is_flash_attn_2_available():
|
|
return "flash_attention_2"
|
|
except ImportError:
|
|
pass
|
|
|
|
# Direct import check for flash_attn 2.x
|
|
try:
|
|
import flash_attn # noqa: F401
|
|
|
|
version = getattr(flash_attn, "__version__", "0.0.0")
|
|
major = int(version.split(".")[0])
|
|
if major >= 2:
|
|
return "flash_attention_2"
|
|
except (ImportError, ValueError, IndexError):
|
|
pass
|
|
|
|
return None
|
|
|
|
|
|
def get_flash_attn_version() -> str | None:
|
|
"""Return the installed flash-attn package version, or None."""
|
|
try:
|
|
import flash_attn
|
|
|
|
return getattr(flash_attn, "__version__", "unknown")
|
|
except ImportError:
|
|
return None
|
|
|
|
|
|
def get_attn_implementation(use_flash_attn: bool, device: str) -> str | None:
|
|
"""Get the best attention implementation to use.
|
|
|
|
Args:
|
|
use_flash_attn: Whether FlashAttention is requested in config.
|
|
device: Training device (cuda/cpu/mps).
|
|
|
|
Returns:
|
|
Attention implementation string for from_pretrained(), or None for default.
|
|
"""
|
|
if not use_flash_attn:
|
|
return None
|
|
|
|
if device != "cuda":
|
|
return None
|
|
|
|
return check_flash_attn_available()
|
|
|
|
|
|
def validate_flash_attn_config(
|
|
use_flash_attn: bool, backend: str, device: str,
|
|
) -> list[str]:
|
|
"""Validate FlashAttention configuration and return error messages.
|
|
|
|
Args:
|
|
use_flash_attn: Whether FlashAttention is requested.
|
|
backend: Training backend (transformers/unsloth).
|
|
device: Training device (cuda/cpu/mps).
|
|
|
|
Returns:
|
|
List of error messages. Empty list means valid.
|
|
"""
|
|
errors: list[str] = []
|
|
|
|
if not use_flash_attn:
|
|
return errors
|
|
|
|
# Unsloth handles FlashAttention internally — no validation needed
|
|
if backend == "unsloth":
|
|
return errors
|
|
|
|
if device != "cuda":
|
|
errors.append(
|
|
"FlashAttention requires CUDA. "
|
|
f"Current device: {device}."
|
|
)
|
|
|
|
if device == "cuda" and check_flash_attn_available() is None:
|
|
errors.append(
|
|
"FlashAttention is not available. "
|
|
"Install it with: pip install flash-attn --no-build-isolation"
|
|
)
|
|
|
|
return errors
|