Nix 中的 CUDA 与机器学习环境：cudaPackages 与 PyTorch

引言

GPU 环境是「依赖地狱」的经典样本：驱动版本、CUDA 运行时、cuDNN、框架（PyTorch/TF）四者必须彼此兼容，而它们由四个不同的团队按不同的节奏发布。Nix 的价值在于把这套矩阵变成可声明、可复现的表达式——同一个 shell 在任何机器上拉到同一套 CUDA。

本文从三方版本约束讲起，说清 nixpkgs 的 cudaPackages 布局、驱动与 cudatoolkit 的关系、如何搭一个能跑 PyTorch 的 GPU 开发 shell，以及版本匹配、容器镜像与编译加速的实践。

前置：Shell 开发环境、语言生态打包、构建 Docker 镜像。

目录

• 1. GPU 环境为什么难：三方版本约束
• 2. cudaPackages 与 nixpkgs 的 CUDA 布局
• 3. 驱动与 cudatoolkit：版本匹配
• 4. GPU 开发 shell 实战
• 5. PyTorch 打包与 Python 生态
• 6. 版本匹配矩阵与常见错配
• 7. 容器镜像与部署
• 8. 编译加速与多 GPU
• 9. 常见坑与排错
• 10. 速查表与一句话记忆
• 延伸阅读

1. GPU 环境为什么难：三方版本约束

1.1 四个层次，四个发布节奏

1.2 传统做法的痛点与 Nix 的解法

传统做法（pip/conda）的痛点是：pip install torch 会带一份自带 CUDA 运行时的 nvidia-* wheel，与系统驱动/系统 CUDA 可能冲突且每个项目各装一份；conda 能缓解但环境不可复现。Nix 的解法：

# 1) 把「运行时 + 库」全部作为 Nix 依赖声明（不进系统全局）
# 2) 只把「驱动」留给系统（NixOS 上有 nvidia 模块，非 NixOS 用宿主驱动）
# 3) 用 shell.nix/flake 固定整套组合 → 可复现、可回滚

记忆：GPU 环境的难点是「驱动/运行时/库/框架」四层各自发布、必须两两兼容；Nix 的解法是把运行时与库全部声明为依赖（不进系统），只把驱动留给宿主系统，从而让整套组合可复现可回滚。

2. cudaPackages 与 nixpkgs 的 CUDA 布局

2.1 cudaPackages 是什么

nixpkgs 把整套 CUDA 工具链组织成一个包集 cudaPackages，可按 CUDA 主版本选择：

{ pkgs, ... }:
{
  # 默认取 nixpkgs 的默认 CUDA 版本
  cuda = pkgs.cudaPackages;
  # 显式选版本（属性名随 nixpkgs 版本变化）
  cuda_12_4 = pkgs.cudaPackages_12_4;
}

2.2 常用成员

2.3 config 层面的开关

{
  nixpkgs.config = {
    # 允许不自由许可（CUDA 属于此类）
    allowUnfree = true;
    cudaSupport = true;
    # 指定 CUDA 版本（若 nixpkgs 支持该开关）
    cudaVersion = "12.4";
  };
}

2.4 为什么必须 allowUnfree

CUDA 与 cuDNN 的许可是「不自由」的，nixpkgs 默认不构建它们。忘记 allowUnfree 会得到「包不存在」或「被许可拒绝」的错误。

记忆：nixpkgs 把 CUDA 组织成可按版本选择的包集 cudaPackages（含 cudatoolkit/cudnn/cublas 等）；用前必须在 nixpkgs.config 里开 allowUnfree = true（CUDA 许可不自由），并用 cudaSupport/cudaVersion 控制版本。

3. 驱动与 cudatoolkit：版本匹配

3.1 驱动与运行时的关系

# 驱动：由宿主内核加载（NixOS 上 nvidia 模块，其他发行版用系统驱动）
# 运行时：Nix 提供（cudatoolkit / cuda_cudart）
# 规则：驱动版本 >= 运行时要求的最低驱动版本

NVIDIA 的向后兼容保证：新版驱动能跑旧版 CUDA 运行时；反之不行。所以「驱动尽量新、运行时按框架要求选」是最稳的策略。

3.2 NixOS 上的驱动声明

{
  # 启用专有驱动
  services.xserver.videoDrivers = [ "nvidia" ];
  hardware.nvidia = {
    modesetting.enable = true;
    open = false;              # 是否用开源内核模块（新卡可 true）
    package = config.boot.kernelPackages.nvidiaPackages.stable;
  };
  # 容器里用 GPU 需要
  hardware.nvidia-container-toolkit.enable = true;
}

3.3 非 NixOS 上只用 Nix 管运行时

{ pkgs, ... }:
{
  # 不装驱动，只用 Nix 提供 CUDA 运行时
  # 前提：宿主已装好兼容的 NVIDIA 驱动
  buildInputs = [ pkgs.cudaPackages.cudatoolkit ];
}

3.4 检查匹配

nvidia-smi                        # 看驱动版本与支持的 CUDA 版本
nvcc --version                    # 看 Nix 提供的运行时版本
# 驱动支持的上限必须 >= 运行时版本

记忆：驱动归宿主（NixOS 用 nvidia 模块）、运行时归 Nix；NVIDIA 保证「新驱动跑旧运行时」，所以策略是驱动尽量新、运行时按框架要求选；nvidia-smi 看驱动上限、nvcc --version 看运行时版本，前者必须 >= 后者。

4. GPU 开发 shell 实战

4.1 一份可用的 shell.nix

{ pkgs ? import <nixpkgs> {
    config.allowUnfree = true;
    config.cudaSupport = true;
  }
}:
let
  cuda = pkgs.cudaPackages;
in
pkgs.mkShell {
  buildInputs = [
    cuda.cudatoolkit
    cuda.cudnn
    cuda.libcublas
    pkgs.python311
    pkgs.python311Packages.numpy
    pkgs.python311Packages.torch-bin   # 见第 5 节
  ];
  # 让编译期能找到 CUDA
  shellHook = ''
    export CUDA_PATH=${cuda.cudatoolkit}
    export LD_LIBRARY_PATH=${cuda.cudatoolkit}/lib:${cuda.cudnn}/lib:$LD_LIBRARY_PATH
    echo "CUDA ready: $(nvcc --version | tail -1)"
  '';
}

4.2 flake 版本与验证

flake 里同样用 import nixpkgs { config = { allowUnfree = true; cudaSupport = true; }; }，把 mkShell 挂到 devShells.<system>.default。验证用：

nix develop          # 或 nix-shell
python -c "import torch; print(torch.cuda.is_available(), torch.version.cuda)"

记忆：GPU 开发 shell 的核心是三件事——allowUnfree + cudaSupport、buildInputs 里放 cudatoolkit/cudnn、shellHook 导出 CUDA_PATH 与 LD_LIBRARY_PATH；验证用 torch.cuda.is_available()。

5. PyTorch 打包与 Python 生态

5.1 torch 与 torch-bin 的区别

日常用 torch-bin；只有在需要改编译选项或打补丁时才用 torch。

5.2 与 poetry2nix 结合

{
  # 用 poetry2nix 把 pyproject.toml 转成 Nix 依赖
  # 注意：torch 需在 overrides 里映射到 torch-bin，否则会触发源码构建
  python311 = pkgs.python311.override {
    packageOverrides = pyfinal: pyprev: {
      torch = pyprev.torch-bin;
    };
  };
}

5.3 常见 Python GPU 库与 pip 的陷阱

常用库直接用 pkgs.python311Packages 里的 torch-bin/torchvision-bin/transformers/accelerate/datasets。

# ☐ pip 装 torch 会带入一份 nvidia-* wheel 的 CUDA 运行时
# ☐ 与 Nix 提供的运行时可能冲突（两份 libcudart）
# ☐ 且该环境不可复现——下次装可能解析到不同版本

记忆：PyTorch 在 nixpkgs 里有 torch（源码，数小时）与 torch-bin（预编译，秒级）两条路——日常用 torch-bin，配合 poetry2nix 时用 packageOverrides 把 torch 映射到 torch-bin；绝不用 pip 装 torch（会带入第二份 CUDA 运行时且不可复现）。

6. 版本匹配矩阵与常见错配

6.1 一份参考矩阵

（具体数值以官方发布矩阵为准，此处示意结构。）

6.2 三类错配

# 1) 驱动太旧：nvidia-smi 显示的 CUDA 上限 < 运行时版本
#    → 症状：CUDA error: no kernel image is available / driver too old
# 2) cuDNN 与 CUDA 主版本不符：cuDNN 9 需 CUDA 12
#    → 症状：导入即崩、找不到 libcudnn.so.9
# 3) 框架与 CUDA 主版本不符：torch 编译时的 CUDA ≠ 运行时
#    → 症状：torch.cuda.is_available() 为 False

6.3 定位方法

用 nvidia-smi | head -4（驱动与支持上限）、nvcc --version | tail -2（运行时）、python -c "import torch; print(torch.version.cuda, torch.backends.cudnn.version())"（框架编译期 CUDA 与 cuDNN）三方对照，三者应构成一条兼容链。

记忆：GPU 错配三态——驱动太旧（上限低于运行时）、cuDNN 与 CUDA 主版本不符、框架编译期 CUDA 与运行时不符；定位靠「nvidia-smi + nvcc –version + torch.version.cuda」三方对照。

7. 容器镜像与部署

7.1 用 dockerTools 构建 GPU 镜像

{ pkgs, ... }:
{
  # 只放 CUDA 运行时与库，驱动由宿主注入
  image = pkgs.dockerTools.buildLayeredImage {
    name = "ml-app";
    contents = [ pkgs.cudaPackages.cudatoolkit pkgs.python311Packages.torch-bin ];
  };
}

7.2 关键原则

# ☐ 镜像里不要放驱动——用 --gpus all 由宿主注入
# ☐ 用 nvidia-container-toolkit 提供 libcuda.so
# ☐ 镜像只带「运行时 + 库 + 应用」，保持层数少

7.3 运行

docker run --rm --gpus all ml-app python -c "import torch; print(torch.cuda.is_available())"

7.4 与官方镜像对比

记忆：GPU 容器镜像的原则是「驱动不进镜像」——用 nvidia-container-toolkit + --gpus all 由宿主注入 libcuda；Nix 用 dockerTools 只打包运行时与库，比官方镜像更可复现、层更少。

8. 编译加速与多 GPU

8.1 cudaCapabilities：只编译需要的架构

{
  nixpkgs.config.cudaCapabilities = [ "8.0" "8.6" "9.0" ];
  # 只编译 Ampere/Ada/Hopper 的 SASS，跳过其他
  # 减少编译时间与产物体积
}

默认会编译一堆 sm_XX 目标，按实际显卡裁剪可显著加速构建。

8.2 多 GPU 与 NCCL

{ pkgs, ... }:
{
  buildInputs = [
    pkgs.cudaPackages.nccl       # 多卡通信
    pkgs.cudaPackages.cudatoolkit
  ];
}

# 指定使用的显卡
CUDA_VISIBLE_DEVICES=0,1 python train.py

8.3 监控与调优

监控用 nvidia-smi dmon -s u（实时利用率）、nvidia-smi --query-gpu=memory.used --format=csv（显存），性能剖析用 nsys profile python train.py。

记忆：编译加速的关键是 cudaCapabilities 只保留实际显卡的架构（默认编一堆 sm_XX 很浪费）；多卡用 cudaPackages.nccl 与 CUDA_VISIBLE_DEVICES 控制可见性，调优用 nvidia-smi dmon 与 nsys。

9. 常见坑与排错

9.1 高频坑

9.2 排错清单

# ☐ 先确认 allowUnfree 与 cudaSupport 都开了
# ☐ nvidia-smi 的 CUDA 上限 >= nvcc --version 的版本
# ☐ LD_LIBRARY_PATH 是否包含 cudatoolkit/lib 与 cudnn/lib
# ☐ 是否误用了 pip 装的 torch（which python 确认在 Nix 环境里）
# ☐ 容器场景确认 --gpus all 与 runtime 配置

9.3 快速自检脚本

一行跑完三方自检：nvidia-smi | head -4; nvcc --version | tail -2; python -c "import torch; print(torch.cuda.is_available(), torch.version.cuda)"。

记忆：CUDA 排错顺序——先确认 allowUnfree/cudaSupport、再对三方版本、再看 LD_LIBRARY_PATH、最后确认没混入 pip 的 torch；容器场景检查 –gpus all 与 nvidia runtime。

10. 速查表与一句话记忆

一句话记忆：GPU 环境难在「驱动/运行时/库/框架」四层各自发布、两两兼容；Nix 的解法是把运行时与库全部声明为依赖（cudaPackages.cudatoolkit/cudnn/...，按版本选包集），只把驱动留给宿主（NixOS 用 nvidia 模块、容器用 nvidia-container-toolkit + --gpus all）；用前必须开 allowUnfree（CUDA 许可不自由）；框架用 torch-bin 而非源码 torch（否则编译数小时），配合 poetry2nix 时用 packageOverrides 映射，绝不用 pip 装 torch（会带入第二份 CUDA 运行时）；shellHook 里导出 CUDA_PATH 与 LD_LIBRARY_PATH 是必备动作；编译加速靠 cudaCapabilities 只保留实际显卡架构；排错顺序是「先确认开关、再对三方版本、再看库路径、最后查是否混入 pip 包」。

延伸阅读

• Shell 开发环境与 direnv
• Nix 语言生态打包
• Nix 构建 Docker 镜像
• 可复现开发环境
• AI 与机器学习专题 — 模型训练与推理实践
• NixOS 维基：CUDA