2.7.7 与 PyTorch 配合使用
原文链接: https://docs.astral.sh/uv/guides/integration/pytorch/
2.7.7 与 PyTorch 配合使用
PyTorch 生态是深度学习研究和开发的热门选择。你可以用 uv 在不同 Python 版本和环境中管理 PyTorch 项目与 PyTorch 依赖，甚至可以控制加速器的选择（例如仅 CPU 还是 CUDA）。
注意
本指南中的部分特性需要 uv 0.5.3 或更高版本。我们建议在配置 PyTorch 之前先升级。
安装 PyTorch
从打包角度看，PyTorch 有几个不同寻常的特点：
因此，所需的打包配置会同时取决于你需要支持的平台和你希望启用的加速器。
首先看下面的（默认）配置，它可以通过依次运行 uv init --python 3.14 和 uv add torch torchvision 生成。
在这种情况下，PyTorch 会从 PyPI 安装，PyPI 为 Windows 和 macOS 提供仅 CPU 的 wheel，为 Linux 提供 GPU 加速 wheel（截至 PyTorch 2.11.0，面向 CUDA 13.0）：
1
2
3
4
5
6
7
8
| [project]
name = "project"
version = "0.1.0"
requires-python = ">=3.14"
dependencies = [
"torch>=2.11.0",
"torchvision>=0.26.0",
]
|
对于希望在 Windows 和 macOS 上使用 CPU 构建、在 Linux 上使用 CUDA 构建的项目，这是一个有效的配置。但如果你需要支持其他平台或加速器，就需要相应配置项目。
使用 PyTorch 索引
某些情况下，你可能希望在所有平台上使用同一种 PyTorch 变体。例如，你可能希望在 Linux 上也使用仅 CPU 的构建。
这种情况下，第一步是把相关的 PyTorch 索引加入 pyproject.toml：
1
2
3
4
| [[tool.uv.index]]
name = "pytorch-cpu"
url = "https://download.pytorch.org/whl/cpu"
explicit = true
|
1
2
3
4
| [[tool.uv.index]]
name = "pytorch-cu118"
url = "https://download.pytorch.org/whl/cu118"
explicit = true
|
1
2
3
4
| [[tool.uv.index]]
name = "pytorch-cu126"
url = "https://download.pytorch.org/whl/cu126"
explicit = true
|
1
2
3
4
| [[tool.uv.index]]
name = "pytorch-cu128"
url = "https://download.pytorch.org/whl/cu128"
explicit = true
|
1
2
3
4
| [[tool.uv.index]]
name = "pytorch-cu130"
url = "https://download.pytorch.org/whl/cu130"
explicit = true
|
1
2
3
4
| [[tool.uv.index]]
name = "pytorch-rocm"
url = "https://download.pytorch.org/whl/rocm7.2"
explicit = true
|
1
2
3
4
| [[tool.uv.index]]
name = "pytorch-xpu"
url = "https://download.pytorch.org/whl/xpu"
explicit = true
|
我们建议使用 explicit = true，以确保该索引仅用于 torch、torchvision 以及其他 PyTorch 相关包，而不会用于 jinja2 这类通用依赖 —— 它们应继续从默认索引（PyPI）获取。
接下来，更新 pyproject.toml，把 torch 和 torchvision 指向所需的索引：
1
2
3
4
5
6
7
| [tool.uv.sources]
torch = [
{ index = "pytorch-cpu" },
]
torchvision = [
{ index = "pytorch-cpu" },
]
|
PyTorch 不为 macOS 发布 CUDA 构建。因此我们用 sys_platform 作为条件，让 uv 在 Linux 和 Windows 上使用 PyTorch 索引，而在 macOS 上回退到 PyPI：
1
2
3
4
5
6
7
| [tool.uv.sources]
torch = [
{ index = "pytorch-cu118", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
]
torchvision = [
{ index = "pytorch-cu118", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
]
|
PyTorch 不为 macOS 发布 CUDA 构建。因此我们用 sys_platform 作为条件，让 uv 只在 Linux 和 Windows 上使用 PyTorch 索引，而在 macOS 上回退到 PyPI：
1
2
3
4
5
6
7
| [tool.uv.sources]
torch = [
{ index = "pytorch-cu126", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
]
torchvision = [
{ index = "pytorch-cu126", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
]
|
PyTorch 不为 macOS 发布 CUDA 构建。因此我们用 sys_platform 作为条件，让 uv 只在 Linux 和 Windows 上使用 PyTorch 索引，而在 macOS 上回退到 PyPI：
1
2
3
4
5
6
7
| [tool.uv.sources]
torch = [
{ index = "pytorch-cu128", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
]
torchvision = [
{ index = "pytorch-cu128", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
]
|
PyTorch 不为 macOS 发布 CUDA 构建。因此我们用 sys_platform 作为条件，让 uv 只在 Linux 和 Windows 上使用 PyTorch 索引，而在 macOS 上回退到 PyPI：
1
2
3
4
5
6
7
| [tool.uv.sources]
torch = [
{ index = "pytorch-cu130", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
]
torchvision = [
{ index = "pytorch-cu130", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
]
|
PyTorch 不为 macOS 或 Windows 发布 ROCm 构建。因此我们用 sys_platform 作为条件，让 uv 只在 Linux 上使用 PyTorch 索引，而在 macOS 和 Windows 上回退到 PyPI：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
| [tool.uv.sources]
torch = [
{ index = "pytorch-rocm", marker = "sys_platform == 'linux'" },
]
torchvision = [
{ index = "pytorch-rocm", marker = "sys_platform == 'linux'" },
]
# ROCm 支持依赖两个 Triton 包，它们也应从 PyTorch 索引安装
# （并包含在 `project.dependencies` 中）。
pytorch-triton-rocm = [
{ index = "pytorch-rocm", marker = "sys_platform == 'linux'" },
]
triton-rocm = [
{ index = "pytorch-rocm", marker = "sys_platform == 'linux'" },
]
|
PyTorch 不为 macOS 发布 Intel GPU 构建。因此我们用 sys_platform 作为条件，让 uv 只在 Linux 和 Windows 上使用 PyTorch 索引，而在 macOS 上回退到 PyPI：
1
2
3
4
5
6
7
8
9
10
11
12
| [tool.uv.sources]
torch = [
{ index = "pytorch-xpu", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
]
torchvision = [
{ index = "pytorch-xpu", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
]
# Intel GPU 支持依赖 `triton-xpu`，它也应从 PyTorch 索引安装
# （并包含在 `project.dependencies` 中）。
triton-xpu = [
{ index = "pytorch-xpu", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
]
|
作为一个完整示例，下面的项目会在所有平台上使用 PyTorch 的仅 CPU 构建：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
| [project]
name = "project"
version = "0.1.0"
requires-python = ">=3.14.0"
dependencies = [
"torch>=2.11.0",
"torchvision>=0.26.0",
]
[tool.uv.sources]
torch = [
{ index = "pytorch-cpu" },
]
torchvision = [
{ index = "pytorch-cpu" },
]
[[tool.uv.index]]
name = "pytorch-cpu"
url = "https://download.pytorch.org/whl/cpu"
explicit = true
|
用环境标记配置加速器
某些情况下，你可能希望在一种环境（例如 macOS 和 Windows）中使用仅 CPU 构建，而在另一种环境（例如 Linux）中使用 CUDA 构建。
借助 tool.uv.sources，你可以用环境标记为每个平台指定所需的索引。例如，下面的配置会在 Linux 上使用 PyTorch 的 CUDA 构建，在其他所有平台（例如 macOS 和 Windows）上使用仅 CPU 构建：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
| [project]
name = "project"
version = "0.1.0"
requires-python = ">=3.14.0"
dependencies = [
"torch>=2.11.0",
"torchvision>=0.26.0",
]
[tool.uv.sources]
torch = [
{ index = "pytorch-cpu", marker = "sys_platform != 'linux'" },
{ index = "pytorch-cu130", marker = "sys_platform == 'linux'" },
]
torchvision = [
{ index = "pytorch-cpu", marker = "sys_platform != 'linux'" },
{ index = "pytorch-cu130", marker = "sys_platform == 'linux'" },
]
[[tool.uv.index]]
name = "pytorch-cpu"
url = "https://download.pytorch.org/whl/cpu"
explicit = true
[[tool.uv.index]]
name = "pytorch-cu130"
url = "https://download.pytorch.org/whl/cu130"
explicit = true
|
类似地，下面的配置会在 Linux 上使用 PyTorch 的 AMD GPU 构建，在 Windows 和 macOS 上使用仅 CPU 构建（通过回退到 PyPI）：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
| [project]
name = "project"
version = "0.1.0"
requires-python = ">=3.14.0"
dependencies = [
"torch>=2.11.0",
"torchvision>=0.26.0",
"pytorch-triton-rocm>=3.5.1 ; sys_platform == 'linux'",
"triton-rocm>=3.6.0 ; sys_platform == 'linux'",
]
[tool.uv.sources]
torch = [
{ index = "pytorch-rocm", marker = "sys_platform == 'linux'" },
]
torchvision = [
{ index = "pytorch-rocm", marker = "sys_platform == 'linux'" },
]
pytorch-triton-rocm = [
{ index = "pytorch-rocm", marker = "sys_platform == 'linux'" },
]
triton-rocm = [
{ index = "pytorch-rocm", marker = "sys_platform == 'linux'" },
]
[[tool.uv.index]]
name = "pytorch-rocm"
url = "https://download.pytorch.org/whl/rocm7.2"
explicit = true
|
或者，针对 Intel GPU 构建：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
| [project]
name = "project"
version = "0.1.0"
requires-python = ">=3.14.0"
dependencies = [
"torch>=2.11.0",
"torchvision>=0.26.0",
"triton-xpu>=3.7.0 ; sys_platform == 'win32' or sys_platform == 'linux'",
]
[tool.uv.sources]
torch = [
{ index = "pytorch-xpu", marker = "sys_platform == 'win32' or sys_platform == 'linux'" },
]
torchvision = [
{ index = "pytorch-xpu", marker = "sys_platform == 'win32' or sys_platform == 'linux'" },
]
triton-xpu = [
{ index = "pytorch-xpu", marker = "sys_platform == 'win32' or sys_platform == 'linux'" },
]
[[tool.uv.index]]
name = "pytorch-xpu"
url = "https://download.pytorch.org/whl/xpu"
explicit = true
|
用可选依赖配置加速器
某些情况下，你可能希望有时使用仅 CPU 构建，有时使用 CUDA 构建，并通过用户提供的 extra 切换（例如 uv sync --extra cpu 与 uv sync --extra cu130）。
借助 tool.uv.sources，你可以用 extra 标记为每个启用的 extra 指定所需索引。例如，下面的配置会在 uv sync --extra cpu 时使用 PyTorch 的仅 CPU 构建，在 uv sync --extra cu130 时使用 CUDA 构建：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
| [project]
name = "project"
version = "0.1.0"
requires-python = ">=3.14.0"
dependencies = []
[project.optional-dependencies]
cpu = [
"torch>=2.11.0",
"torchvision>=0.26.0",
]
cu130 = [
"torch>=2.11.0",
"torchvision>=0.26.0",
]
[tool.uv]
conflicts = [
[
{ extra = "cpu" },
{ extra = "cu130" },
],
]
[tool.uv.sources]
torch = [
{ index = "pytorch-cpu", extra = "cpu" },
{ index = "pytorch-cu130", extra = "cu130" },
]
torchvision = [
{ index = "pytorch-cpu", extra = "cpu" },
{ index = "pytorch-cu130", extra = "cu130" },
]
[[tool.uv.index]]
name = "pytorch-cpu"
url = "https://download.pytorch.org/whl/cpu"
explicit = true
[[tool.uv.index]]
name = "pytorch-cu130"
url = "https://download.pytorch.org/whl/cu130"
explicit = true
|
注意
由于 macOS 上没有 GPU 加速构建，启用 cu130 extra 时上述配置在 macOS 上会安装失败。
安装支持 GPU 的 PyTorch 扩展
PyTorch 生态中的许多包都包含针对特定 CUDA 与 PyTorch 版本组合编译的 GPU 扩展。从源码构建这些包通常需要 CUDA 开发工具包以及额外的构建配置。
Astral GPU 索引为 flash-attn、deepspeed、deep-gemm、torch-scatter 和 vllm 等包提供预构建 wheel，覆盖多种 Python、CUDA 和 PyTorch 版本。
要从 CUDA 12.8 的索引安装 flash-attn，请运行：
1
| $ uv add flash-attn --index astral-cu128=https://wheels.astral.sh/simple/cu128/
|
该命令会把 flash-attn 加入项目依赖、配置 Astral GPU 索引，并把 flash-attn 固定到该索引。
与 PyTorch 索引一样，设置 explicit = true 可以把 Astral GPU 索引限制为仅用于显式固定到它的包：
1
2
3
4
5
6
7
8
| # pyproject.toml
[tool.uv.sources]
flash-attn = { index = "astral-cu128" }
[[tool.uv.index]]
name = "astral-cu128"
url = "https://wheels.astral.sh/simple/cu128/"
explicit = true
|
每个 Astral GPU 索引面向特定的 CUDA 版本，其 wheel 也是为特定 PyTorch 版本构建的。例如，本地版本为 +cu.12.8.torch.2.11 的 wheel 是为 CUDA 12.8 和 PyTorch 2.11 构建的。请选择与你的 Python 版本、平台、CUDA 版本和 PyTorch 安装相匹配的索引与 wheel。
可用的包、CUDA 版本和 PyTorch 版本请参阅 Astral GPU 索引。
uv pip 接口
虽然上面的示例聚焦于 uv 的项目接口（uv lock、uv sync、uv run 等），PyTorch 也可以通过 uv pip 接口安装。
PyTorch 本身提供了专用界面，用于确定针对给定目标配置应运行哪条 pip 命令。例如，你可以这样在 Linux 上安装稳定的、仅 CPU 的 PyTorch：
1
| $ pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
|
要在 uv 中使用同样的工作流，请把 pip3 换成 uv pip：
1
| $ uv pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
|
自动选择后端
uv 支持通过 --torch-backend=auto 命令行参数（或 UV_TORCH_BACKEND=auto 环境变量）自动选择合适的 PyTorch 索引，例如：
1
2
3
4
5
| $ # 使用命令行参数。
$ uv pip install torch --torch-backend=auto
$ # 使用环境变量。
$ UV_TORCH_BACKEND=auto uv pip install torch
|
启用后，uv 会查询已安装的 CUDA 驱动、AMD GPU 版本以及是否存在 Intel GPU，然后为所有相关包（例如 torch、torchvision 等）使用兼容性最好的 PyTorch 索引。如果没有找到这类 GPU，uv 会回退到仅 CPU 索引。对于 PyTorch 生态之外的包，uv 会继续遵循现有的索引配置。
你也可以用 --torch-backend=cu130（或 UV_TORCH_BACKEND=cu130）选择特定后端（例如 CUDA 13.0）：
1
2
3
4
5
| $ # 使用命令行参数。
$ uv pip install torch torchvision --torch-backend=cu130
$ # 使用环境变量。
$ UV_TORCH_BACKEND=cu130 uv pip install torch torchvision
|
目前 --torch-backend 仅在 uv pip 接口中可用。