标签 flash-attn 下的文章

如果你拥有一个纯净的 Ubuntu 22.04 环境,想要一次性配置好“AI 全家桶”开发环境,最稳妥、最标准的方法是遵循 NVIDIA 网络仓库安装流程

这套流程不仅会安装编译器(nvcc),还会安装驱动(Driver)、数学库(cuBLAS/cuDNN)以及开发工具,确保你后续安装 flash-attnvLLMDeepSpeed 时一马平川。

第一步:基础环境准备

首先确保系统有基础的下载和编译工具:

sudo apt update
sudo apt install -y build-essential charts-common wget curl git

第二步:配置 NVIDIA 官方仓库 (最关键)

通过官方提供的 cuda-keyring,让 apt 能够获取到所有 AI 相关的专用包。

# 1. 下载官方密钥配置包
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb

# 2. 安装密钥配置包
sudo dpkg -i cuda-keyring_1.1-1_all.deb

# 3. 更新仓库索引
sudo apt update

第三步:一键安装 CUDA + 驱动 + cuDNN

在纯净环境下,你可以直接安装 cuda 完整包,它会自动帮你匹配并安装最合适的显卡驱动。

# 安装 CUDA 完整版(包含驱动、Toolkit、示例、NVCC 等)
sudo apt -y install cuda

# 安装深度学习必备的 cuDNN 库
sudo apt -y install libcudnn9-cuda-12

注:由于我们正处在 2026 年,如果系统提示找不到 libcudnn9,可以尝试 sudo apt install libcudnn-cuda-12 来获取最新的兼容版本。


第四步:永久配置环境变量

安装在 /usr/local/cuda 的工具如果不加进 PATH,系统是识别不到 nvcc 的。

# 将配置写入当前用户的 .zshrc (如果是 bash 则改为 .bashrc)
echo 'export CUDA_HOME=/usr/local/cuda' >> ~/.zshrc
echo 'export PATH=$CUDA_HOME/bin:$PATH' >> ~/.zshrc
echo 'export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH' >> ~/.zshrc

# 立即生效
source ~/.zshrc

第五步:验证安装

运行以下两条命令,如果都能看到信息,说明你的 AI 塔基已经打好了:

  1. nvidia-smi:检查显卡驱动和硬件状态。
  2. nvcc --version:检查 CUDA 编译器是否就绪。

进阶:针对 AI 开发的额外补充

为了让你以后跑模型更顺滑,建议顺手装上这几个库:

  1. NCCL(多显卡通信库,分布式训练必装):

    sudo apt install libnccl2 libnccl-dev
    
  2. NVML(监控库):通常已随驱动安装。

总结你的“一键式”清单:

  1. dpkg -i cuda-keyring... (换源)
  2. apt install cuda (装驱动+工具链)
  3. apt install libcudnn... (装深度学习加速库)
  4. export PATH... (设环境变量)

当你完成这些后,再去你的虚拟环境中安装 flash-attn,它就会自动找到 nvcc 并开始愉快的编译了。需要我帮你写一个自动执行这些步骤的 .sh 脚本吗?