Flash-Attention-2安装和后续排错

在https://zhuanlan.zhihu.com/p/1897638897388873200基础上。

前言

Flash-Attention的安装其实并没有那么复杂，网上的帖子有很多，但不够简明扼要。亲测按照以下步骤，大概20min之后就可以安装成功。

要求

CUDA >= 12.0 ；torch >=2.2
仅支持Ampere, Ada, or Hopper GPUs (比如 A100, RTX 3090, RTX 4090, H100)。不支持V100
建议cuda和torch的版本相匹配，参考
pytorch-version
pytorch.org/get-started/previous-versions/
首先检查本地python、torch、cuda的版本

python --version &&
python -c "import torch; print(torch.__version__); print(torch.cuda.is_available())" &&
nvcc -V

执行以上命令，输出：python版本3.10 torch版本2.5 cuda版本12.1

Python 3.10.14
2.5.1+cu124
True
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2023 NVIDIA Corporation
Built on Mon_Apr__3_17:16:06_PDT_2023
Cuda compilation tools, release 12.1, V12.1.105
Build cuda_12.1.r12.1/compiler.32688072_0

这里最好是cuda和torch的版本相匹配(但其实如果不匹配也没有问题)，不匹配的话pip安装对应的版本即可。
主要是python和cuda的版本要匹配 whl

安装

查询好本地的版本后，从这里下载对应的whl

Releases · Dao-AILab/flash-attention
github.com/Dao-AILab/flash-attention/releases
打开之后应该是这样的：

flash_attn：这是包的名称。
2.7.4.post1：这是包的版本号，post1 表示这是版本 2.7.4 的一个后续修订版本。
cu12：表示该包是针对 CUDA 12 版本编译的
torch2.5：表明该包与 PyTorch 2.5 版本兼容。
cxx11abi：指的是 C++ ABI（应用二进制接口），cxx11abiTRUE 表示使用 C++11 ABI 编译，cxx11abiFALSE 则表示不使用 C++11 ABI 编译。不同的编译器和系统可能对 C++ ABI 有不同的支持，因此需要根据自己的环境选择合适的版本。
cp311：代表 Python 解释器的版本，这里是 CPython 3.11。文件名中出现两次 cp311，第一个表示编译该包时使用的 Python 版本，第二个表示该包所兼容的 Python 版本。

下载好对应的whl后，直接pip安装即可

pip install flatsh_atten-xxxxx-xxxxx.whl

使用

Basic

model = AutoModel.from_pretrained(
    "model_path",
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2",
    device_map="auto",
)

解决数据集加载速度

如果出现卡顿，那就是数据集加载速度过慢。

可以考虑

dataloader_num_workers=128,       # 设置多进程读取数据（核心！如果卡顿，调大这个）
dataloader_pin_memory=True,     # 开启锁页内存，加快从 CPU 到 GPU 的拷贝
dataloader_persistent_workers=True, # 保持 worker 进程存活，避免每个 Epoch 重启进程的开销

这几个方面

解决transformers库版本太低，导致无法加载参数的问题

直接强制写入

model_args, data_args, training_args = parser.parse_args_into_dataclasses()
training_args.dataloader_persistent_workers = True

dataloader加载太多文件导致 Too many open files

# 查看当前限制
ulimit -n

# 提升限制
ulimit -n 65535

# 然后再次运行你的训练脚本
python /nvem/projects/gpn/./gpn/msa/train.py ...