用 uv 装 PyTorch + 正确 CUDA 版本（一条命令搞定）

PyTorch 装错 CUDA 版本（CPU-only / 不匹配显卡驱动）是新手最常踩的
第一个坑。uv + PyTorch 官方提供的 index URL 能在一行命令里搞定。

1. 看清你需要的版本

nvidia-smi
# 关注右上角 "CUDA Version: 12.4" —— 这是 driver 支持的最高 CUDA 版本
# PyTorch 实际使用的 runtime 可以等于或更低，但不能更高

PyTorch 现在主要 ship 三种轮子：

cu124 (CUDA 12.4 runtime)
cu121 (CUDA 12.1)
cu118 (CUDA 11.8) — 老显卡 / 老驱动
cpu (无 GPU)

你的驱动支持 12.x 就装 cu121 或 cu124；不确定就 cu121（兼容性最好）。

2. 一条命令装

# uv 新项目
uv init mlproject --python 3.12
cd mlproject

# 直接装最新稳定版
uv add torch torchvision torchaudio --index https://download.pytorch.org/whl/cu124

# 或指定版本
uv add 'torch==2.4.1' 'torchvision==0.19.1' \
  --index https://download.pytorch.org/whl/cu124

--index 让 uv 从 PyTorch 自己的 CDN 拉轮子（PyPI 上的轮子默认是 CPU 版的）。

3. 校验

# check.py
import torch

print(f'torch:        {torch.__version__}')
print(f'cuda built:   {torch.version.cuda}')
print(f'cuda avail:   {torch.cuda.is_available()}')
print(f'device count: {torch.cuda.device_count()}')
if torch.cuda.is_available():
    print(f'device name:  {torch.cuda.get_device_name(0)}')
    print(f'capability:   {torch.cuda.get_device_capability(0)}')

# 做一次实际计算
x = torch.randn(1000, 1000)
if torch.cuda.is_available():
    x = x.cuda()
    print(f'matmul on GPU: {(x @ x).sum().item():.2f}')

uv run python check.py

期望输出：

torch:        2.4.1+cu124
cuda built:   12.4
cuda avail:   True
device count: 1
device name:  NVIDIA GeForce RTX 4090
capability:   (8, 9)
matmul on GPU: ...

+cu124 后缀确认装的是 CUDA wheel；is_available() False 说明 wheel
匹配了但驱动 / NVIDIA 库有问题。

4. 锁版本

ls uv.lock         # 已经写入

CI / 同事直接：

uv sync --frozen --index https://download.pytorch.org/whl/cu124

得到完全一致的环境。

5. 配置 pyproject.toml 让 index 持久化

每次都加 --index 很烦。pyproject.toml：

[project]
name = "mlproject"
version = "0.1.0"
requires-python = ">=3.12"
dependencies = [
    "torch>=2.4",
    "torchvision>=0.19",
]

[[tool.uv.index]]
name = "pytorch-cu124"
url = "https://download.pytorch.org/whl/cu124"
explicit = true

[tool.uv.sources]
torch = { index = "pytorch-cu124" }
torchvision = { index = "pytorch-cu124" }

之后 uv add ... 不再需要 --index。

6. Apple Silicon (M1/M2/M3) 用 MPS

uv add torch torchvision torchaudio
# Mac 上 PyPI 默认轮子已经带 MPS（Metal Performance Shaders）后端

device = 'mps' if torch.backends.mps.is_available() else 'cpu'
x = torch.randn(1000, 1000, device=device)

MPS 不如 CUDA 快但比 CPU 快几倍，免费的不错。

7. CPU-only（开发 / 测试 / CI）

uv add torch --index https://download.pytorch.org/whl/cpu

CI 跑测试一般用 CPU wheel，节省 wheel 大小 + 启动时间。

8. 多 GPU 看见

torch.cuda.device_count()      # 几张卡
# 显式选择
device = torch.device('cuda:0')
model = model.to(device)

多卡训练用 torch.nn.DataParallel 简单粗暴 / DistributedDataParallel
（DDP）才是生产选择。

9. PyTorch 升级

uv lock --upgrade-package torch --index https://download.pytorch.org/whl/cu124

不同 CUDA 版本的 PyTorch 是不同 package（torch+cu121 vs torch+cu124），
直接 lock 升级最稳。

10. Docker 镜像

FROM nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04

# uv 二进制
COPY --from=ghcr.io/astral-sh/uv:latest /uv /usr/local/bin/uv

# Python 3.12
RUN apt update && apt install -y python3.12 python3-pip git

WORKDIR /app
COPY pyproject.toml uv.lock ./
RUN uv sync --frozen --index https://download.pytorch.org/whl/cu124
COPY . .
CMD ["uv", "run", "python", "train.py"]

镜像底层用 NVIDIA 官方 cuda:*-devel，宿主装 nvidia-container-toolkit：

docker run --gpus all -v $(pwd):/app mlproject

踩过的坑

装了 CPU wheel 但代码里 .cuda()：报 "no CUDA-capable device"。
pip / uv 默认从 PyPI 拉 = CPU wheel。必须 --index pytorch。
驱动 CUDA 12.0 但装了 cu124：装得上但 is_available()=False 或运行
时 segfault。升驱动或者降 wheel。
多版本 Python 共存时 wheel 不匹配（cp310 wheel 装到 cp312 venv）：
uv 会自动挑对的，但手动 pip 时容易选错。
想用 torch.compile() 但 wheel 没带：cu118 上没 compile 支持。
cu121+ / cu124 才有。
jupyter notebook 启动时不见 CUDA：notebook 是另一个进程，
确认是用 uv run jupyter 启动（继承了 venv 的 path）。