Spaces:
Sleeping
Sleeping
| """ | |
| Attention-backend selection helper. | |
| Picks the fastest attention implementation available at runtime, with a | |
| safe fallback ladder: | |
| flash_attention_2 (package `flash-attn`) | |
| ↓ not installed / incompatible | |
| sdpa (torch.nn.functional.scaled_dot_product_attention) | |
| ↓ not supported by this model | |
| eager (stock HF implementation, slowest) | |
| Flash-Attn 2 is a big deal for this codebase: | |
| * Training (PPO backward pass): | |
| - Turns attention activation memory from O(T^2) to O(T) per layer. | |
| For B=8, T=500, H=12, 28 layers of bf16 Qwen2 that is a ~1.3 GB | |
| saving on the backward graph, and it scales quadratically with T. | |
| - Fused forward+backward is 1.5-2.5x faster than SDPA on Ampere+. | |
| * Rollouts (KV-cached `.generate()`): | |
| - Each generation step does an incremental attention over the full | |
| KV cache. Flash is faster here too, and its lower memory footprint | |
| lets us keep larger prompts cached. | |
| The helper memoizes its answer so we only probe the `flash_attn` import | |
| once per process. | |
| """ | |
| from __future__ import annotations | |
| import logging | |
| from typing import Optional | |
| logger = logging.getLogger(__name__) | |
| # Module-level cache. Set to None the first time select_attn_implementation | |
| # runs; subsequent calls return the cached string. | |
| _SELECTED: Optional[str] = None | |
| def select_attn_implementation( | |
| prefer: Optional[str] = None, | |
| log_once: bool = True, | |
| ) -> str: | |
| """ | |
| Pick the best attention backend string for | |
| `AutoModel{,ForCausalLM}.from_pretrained(..., attn_implementation=...)`. | |
| Args: | |
| prefer: | |
| If set, try this backend first. Useful for forcing "sdpa" | |
| in environments where flash-attn is installed but broken | |
| (rare, but we've seen it on some vast.ai images). | |
| log_once: | |
| When True, emit one INFO log line the first time we pick, | |
| then be silent on subsequent calls. | |
| Returns: | |
| One of: "flash_attention_2", "sdpa", "eager". | |
| """ | |
| global _SELECTED | |
| if _SELECTED is not None: | |
| return _SELECTED | |
| candidates = [] | |
| if prefer is not None: | |
| candidates.append(prefer) | |
| # Canonical preference order. | |
| for name in ("flash_attention_2", "sdpa", "eager"): | |
| if name not in candidates: | |
| candidates.append(name) | |
| chosen = "eager" | |
| for name in candidates: | |
| if name == "flash_attention_2": | |
| if _flash_attention_2_available(): | |
| chosen = "flash_attention_2" | |
| break | |
| elif name == "sdpa": | |
| # SDPA ships with torch >= 2.0 and is always importable. | |
| # The HF model class may still reject it for non-supported | |
| # architectures, but every modern Llama/Qwen supports it. | |
| chosen = "sdpa" | |
| break | |
| elif name == "eager": | |
| chosen = "eager" | |
| break | |
| _SELECTED = chosen | |
| if log_once: | |
| logger.info( | |
| "Attention backend selected: %s%s", | |
| chosen, | |
| "" if chosen == "flash_attention_2" else | |
| " (flash-attn not available — install `flash-attn` for " | |
| "~1.5-2.5x faster attention and O(T) memory)", | |
| ) | |
| return chosen | |
| def _flash_attention_2_available() -> bool: | |
| """ | |
| Return True iff `flash_attn` is importable and its version is >=2.0. | |
| We don't run a functional test; HF will raise a clear error at | |
| model-load time if the installed build is incompatible with the | |
| model's head dim / dtype, and we'd rather surface that than silently | |
| fall back and waste hours of training at 1x speed. | |
| """ | |
| try: | |
| import flash_attn # noqa: F401 | |
| except Exception: | |
| return False | |
| version = getattr(flash_attn, "__version__", "0.0") | |
| try: | |
| major = int(str(version).split(".", 1)[0]) | |
| except ValueError: | |
| return False | |
| return major >= 2 | |