flash-attn

FlashAttention 加速注意力

补充说明

FlashAttention 通过分块和重计算减少显存访问，提升长上下文训练的吞吐并降低显存占用。flash-attn 是常用实现，支持 Ampere、Hopper 等较新 GPU。

它通常作为 PyTorch 的注意力后端使用，Transformers、vLLM、DeepSpeed 等框架会自动检测。

常用示例

安装：

pip install flash-attn --no-build-isolation
python -c "import flash_attn; print(flash_attn.__version__)"

在 Transformers 中启用：

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2.5-7B-Instruct",
    torch_dtype="bfloat16",
    attn_implementation="flash_attention_2",
    device_map="auto",
)

检查 GPU 架构：

nvidia-smi --query-gpu=name,compute_cap --format=csv

基准测试：

python -m flash_attn.utils.benchmark

易错点与排错

• 编译需要 CUDA toolkit 和匹配的 PyTorch 版本，--no-build-isolation 可避免依赖隔离问题。
• 旧 GPU（如 V100/T4）不支持 FlashAttention 2，会报架构不兼容。
• PyTorch、CUDA、flash-attn 三者版本不匹配，导入时报 undefined symbol。
• 安装耗时很长，建议使用预编译 wheel 或固定版本。
• attn_implementation 写错，Transformers 回退到 eager 但不报错。
• 长上下文收益明显，短序列可能没有提升甚至略慢。
• 与自定义注意力 mask 或位置编码不兼容时，输出可能异常。
• 多卡训练要确保每个进程都安装了正确版本。