记录一下那天晚上的排查过程。pip install flash-attn 卡在编译阶段,一个多小时后报 nvcc 内存不足或编译错误,机器有 64G 内存仍然失败。
刚开始盯着「并行编译进程太多,内存被瞬间吃光」,折腾了半个多小时,排掉了。
原因不止一个,得分开看。源码编译需要大量内存与匹配的 CUDA 工具链。另一种情况是Python 与 torch 版本没有对应的预编译轮子。
还有一种情况,CUDA 版本与 wheel 标签不匹配。
动手时先查是否有匹配当前环境的预编译包
pip install flash-attn --no-build-isolation
限制编译并发以降低内存峰值
MAX_JOBS=2 pip install flash-attn --no-build-isolation
随后确认 torch 与 CUDA 版本组合
python -c "import torch;print(torch.__version__, torch.version.cuda)"
装好后验证能否导入
python -c "import flash_attn;print(flash_attn.__version__)"
命令输出我都留了记录,免得改完记不清原来的值。
改完先别撤,确认这几项。安装成功且导入不报错,训练时开启该注意力实现,显存占用下降、速度提升。
编译型依赖最好在目标环境里一次装好并做镜像,换机器重装往往比调编译参数更省时间。如果同一现象隔几天又出现,说明根因没找到,别当偶发处理。