uv 使用 PyTorch 指南从 CPU 到 CUDA、ROCm 与 Intel GPU 的加速器感知配置【免费下载链接】uvAn extremely fast Python package and project manager, written in Rust.项目地址: https://gitcode.com/GitHub_Trending/uv/uv本文基于 uv 官方文档docs/guides/integration/pytorch.md系统讲解如何用 uv 管理 PyTorch 项目如何按平台与加速器CPU-only、CUDA、ROCm、Intel GPU选择并固定 PyTorch 索引、如何用环境标记与 optional dependencies 表达多环境策略以及--torch-backendauto自动后端选择的底层检测机制。读完后你可以为任意平台组合编写可复制、可锁定的pyproject.toml配置并理解 uv 在源码层面如何决定“给这台机器装哪个 PyTorch”。需要说明的适用前提文档指出文中部分功能要求 uv 版本 0.5.3 或更高建议在配置 PyTorch 前先升级 uv此外--torch-backend选项在 CLI 源码 中被标注为 preview预览特性未来版本中可能发生变化。PyTorch 的打包特性为什么需要特殊配置从打包角度看PyTorch 有几个不寻常的特点引自 docs/guides/integration/pytorch.md大量 wheel 托管在专属索引上而不是 PyPI。安装 PyTorch 通常需要为项目配置 PyTorch 索引。每个加速器CPU-only、CUDA 等对应独立的构建。由于发布和安装时没有标准化的加速器指定机制PyTorch 把加速器编码在本地版本说明符local version specifier中因此版本看起来像2.11.0cpu、2.11.0cu130。不同加速器的构建发布在不同索引上。例如cpu构建发布在https://download.pytorch.org/whl/cpucu130构建发布在https://download.pytorch.org/whl/cu130。因此所需的打包配置取决于你要支持的平台和要启用的加速器。默认配置PyPI 上的 PyTorch最简单的情形由uv init --python 3.14后运行uv add torch torchvision生成。此时 PyTorch 从 PyPI 安装而 PyPI 托管的是 Windows 与 macOS 上的 CPU-only wheel以及 Linux 上的 GPU 加速 wheel截至 PyTorch 2.11.0针对 CUDA 13.0[project] name project version 0.1.0 requires-python 3.14 dependencies [ torch2.11.0, torchvision0.26.0, ]这是一个合法的起点Windows 和 macOS 上使用 CPU 构建Linux 上使用 CUDA 构建。但如果你需要支持其他平台组合或加速器就需要按下面的方式显式配置项目。在项目中指定 PyTorch 索引若想在所有平台统一使用某个 PyTorch 变体例如在 Linux 上也要用 CPU-only 构建第一步是把对应的 PyTorch 索引写入pyproject.toml。官方文档给出的各变体索引如下变体索引名name索引地址urlCPU-onlypytorch-cpuhttps://download.pytorch.org/whl/cpuCUDA 11.8pytorch-cu118https://download.pytorch.org/whl/cu118CUDA 12.6pytorch-cu126https://download.pytorch.org/whl/cu126CUDA 12.8pytorch-cu128https://download.pytorch.org/whl/cu128CUDA 13.0pytorch-cu130https://download.pytorch.org/whl/cu130ROCm 7.2pytorch-rocmhttps://download.pytorch.org/whl/rocm7.2Intel GPU (XPU)pytorch-xpuhttps://download.pytorch.org/whl/xpu以 CUDA 13.0 为例索引声明为[[tool.uv.index]] name pytorch-cu130 url https://download.pytorch.org/whl/cu130 explicit true官方推荐explicit true确保该索引只用于torch、torchvision等 PyTorch 相关包而jinja2这类通用依赖仍从默认索引PyPI解析。第二步用[tool.uv.sources]把torch和torchvision指向目标索引。不同变体需要不同的环境标记marker原因如下CPU-only所有平台都适用不需要 marker[tool.uv.sources] torch [ { index pytorch-cpu }, ] torchvision [ { index pytorch-cpu }, ]CUDA 11.8 / 12.6 / 12.8 / 13.0PyTorch 不为 macOS 发布 CUDA 构建因此用sys_platform限制索引只作用于 Linux 和 WindowsmacOS 回落到 PyPImacOS 的 PyPI wheel 即 CPU 版[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 }, ]ROCm 7.2PyTorch 不为 macOS 或 Windows 发布 ROCm 构建marker 只保留 Linux并且 ROCm 支持依赖两个 Triton 包它们也必须从 PyTorch 索引安装同时列入project.dependencies[tool.uv.sources] torch [ { index pytorch-rocm, marker sys_platform linux }, ] torchvision [ { index pytorch-rocm, marker sys_platform linux }, ] # ROCm support relies on both Triton packages, which should also be installed from # the PyTorch index (and included in project.dependencies). pytorch-triton-rocm [ { index pytorch-rocm, marker sys_platform linux }, ] triton-rocm [ { index pytorch-rocm, marker sys_platform linux }, ]Intel GPUPyTorch 不为 macOS 发布 Intel GPU 构建marker 限制在 Linux 与 WindowsIntel GPU 支持依赖triton-xpu同样要从 PyTorch 索引安装[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 support relies on triton-xpu, which should also be installed from the PyTorch index # (and included in project.dependencies). triton-xpu [ { index pytorch-xpu, marker sys_platform linux or sys_platform win32 }, ]完整示例所有平台使用 CPU-only 构建把索引声明与 sources 组合起来一个完整项目如下[project] name project version 0.1.0 requires-python 3.14.0 dependencies [ torch2.11.0, torchvision0.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-only 构建在 Linux 上用 CUDA 构建。tool.uv.sources允许对同一个包列出多条带 marker 的索引来源uv 会按 marker 在锁定时选择生效的一条。下面的配置让 Linux 使用 CUDA 13.0其他平台使用 CPU-only[project] name project version 0.1.0 requires-python 3.14.0 dependencies [ torch2.11.0, torchvision0.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 使用 AMD GPUROCm 7.2构建Windows 与 macOS 通过回落 PyPI 获得 CPU 构建。注意 ROCm 场景下 Triton 依赖包也必须写入dependencies并钉到 ROCm 索引[project] name project version 0.1.0 requires-python 3.14.0 dependencies [ torch2.11.0, torchvision0.26.0, pytorch-triton-rocm3.5.1 ; sys_platform linux, triton-rocm3.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 构建把 ROCm 相关部分替换为 XPU 即可[project] name project version 0.1.0 requires-python 3.14.0 dependencies [ torch2.11.0, torchvision0.26.0, triton-xpu3.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用可选依赖extra切换加速器有时希望由使用者在同步时决定装哪个变体例如uv sync --extra cpu与uv sync --extra cu130。tool.uv.sources支持 extra marker可为每个启用的 extra 指定不同的索引[project] name project version 0.1.0 requires-python 3.14.0 dependencies [] [project.optional-dependencies] cpu [ torch2.11.0, torchvision0.26.0, ] cu130 [ torch2.11.0, torchvision0.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这里的关键点是[tool.uv].conflicts它声明cpu与cu130两个 extra 互斥避免 uv 在锁定时把两个互相冲突的索引策略同时解析进同一环境。官方文档特别提醒由于 GPU 加速构建在 macOS 上不可用上面的配置在启用cu130extra 时于 macOS 上会安装失败——这正是 expected 的行为也是为什么跨平台团队更推荐上一节的 marker 方案。安装启用 GPU 的 PyTorch 扩展PyTorch 生态中许多包带有针对特定 CUDA 与 PyTorch 版本组合编译的 GPU 扩展。从源码构建这类包通常需要 CUDA 开发工具链和额外的构建配置成本很高。Astral 提供的 GPU 索引Astral GPU indexes为flash-attn、deepspeed、deep-gemm、torch-scatter、vllm等包提供预构建 wheel覆盖多种 Python、CUDA 与 PyTorch 版本组合。要为 CUDA 12.8 安装flash-attn运行$ uv add flash-attn --index astral-cu128https://wheels.astral.sh/simple/cu128/该命令会把flash-attn加入项目依赖、配置 Astral GPU 索引并把flash-attn固定到该索引等价于生成如下配置[tool.uv.sources] flash-attn { index astral-cu128 } [[tool.uv.index]] name astral-cu128 url https://wheels.astral.sh/simple/cu128/ explicit true与 PyTorch 索引一样设置explicit true可将 Astral GPU 索引限制在显式钉到它的包上避免“污染”其他依赖的解析。注意 Astral GPU 索引的 wheel 命名规则每个索引针对特定 CUDA 版本wheel 针对特定 PyTorch 版本构建。例如本地版本为cu.12.8.torch.2.11的 wheel 对应 CUDA 12.8 与 PyTorch 2.11。选择索引和 wheel 时要匹配你自己的 Python 版本、平台、CUDA 版本和已安装的 PyTorch 版本。具体可用包、CUDA 版本与 PyTorch 版本列表可查阅 Astral GPU indexes 站点配置中引用的https://wheels.astral.sh/simple/cu128/即该站点的一个端点。uv pip接口下的 PyTorch 安装前面所有示例都基于 uv 的项目接口uv lock、uv sync、uv run等但 PyTorch 同样可以通过uv pip接口安装。PyTorch 官方提供按目标配置生成 pip 命令的界面。例如在 Linux 上安装稳定版 CPU-only PyTorch 的标准命令是$ pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu用 uv 时只需把pip3换成uv pip$ uv pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu自动后端选择--torch-backendauto除了手工配置索引uv 还支持自动选择适当的 PyTorch 索引通过--torch-backendauto命令行参数或UV_TORCH_BACKENDauto环境变量启用$ # With a command-line argument. $ uv pip install torch --torch-backendauto $ # With an environment variable. $ UV_TORCH_BACKENDauto uv pip install torch启用后uv 会查询已安装的 CUDA 驱动、AMD GPU 版本和 Intel GPU 的存在性然后为所有相关包如torch、torchvision等选用兼容性最好的 PyTorch 索引如果未发现任何 GPU则回落到 CPU-only 索引。对于 PyTorch 生态之外的包uv 仍继续尊重已有的索引配置。也可以直接指定某个后端例如 CUDA 13.0用--torch-backendcu130或UV_TORCH_BACKENDcu130$ # With a command-line argument. $ uv pip install torch torchvision --torch-backendcu130 $ # With an environment variable. $ UV_TORCH_BACKENDcu130 uv pip install torch torchvision按文档说明--torch-backend目前仅在uv pip接口中可用。源码级剖析uv 如何决定 PyTorch 索引uv 的自动后端选择由独立 crate uv-torch 实现其定位就是“根据操作系统和 CUDA 驱动版本决定合适的 PyTorch 索引”实现源自 light-the-torch 项目见 backend.rs 的文件头注释。1. 后端取值范围远大于文档表格。TorchMode 枚举 定义了auto、cpu、CUDA 13.2 到 CUDA 10.0 的各版本cu132、cu130、cu129、cu128、cu126一路到cu80、ROCm 7.2 到 4.0.1 的各版本rocm7.2、rocm7.1、rocm6.4……以及xpu。也就是说CLI 接受的后端值比文档表格列出的那几个更完整文档表格只展示了常用的几个。2.auto模式的硬件探测顺序。探测逻辑在 accelerator.rs 中Accelerator::detect的文档注释列出了严格的查询优先级UV_CUDA_DRIVER_VERSION环境变量覆盖值UV_AMD_GPU_ARCHITECTURE环境变量覆盖值/sys/module/nvidia/versionNVIDIA 驱动版本如550.144.03/proc/driver/nvidia/versionnvidia-smi --query-gpudriver_version --formatcsv,noheaderrocm_agent_enumerator列出 AMD GPU 架构如gfx906/sys/bus/pci/devices按 PCI 厂商号0x8086Intel与显示类码过滤出 Intel GPUWindows 设备树按 PCI 过滤已安装的 Intel 显示适配器。值得注意的是这里检测的是CUDA 驱动版本如550.144.03而非 CUDA 工具包版本如12.8.0——源码注释明确区分了两者见 accelerator.rs。index_urls 方法 随后根据操作系统把驱动版本映射为索引列表Linux 上按驱动版本从高到低筛出所有满足driver_version 最低要求的 CUDA 索引最后总是追加 CPU-only 索引作为兜底Windows 逻辑相同macOS 及其他平台只给出 CPU-only 索引。这解释了为什么文档中反复强调“macOS 回落 PyPI”——自动模式下 macOS 根本不会产生 CUDA 索引。3. “相关包”的边界由源码固定。TorchStrategy::applies_to 硬编码了自动策略作用的包名清单torch、torchvision、torchaudio、torchtext、torchtune、torchao、torchrec、fbgemm-gpu、triton、triton-rocm、triton-xpu、xformers等has_system_dependency 进一步区分哪些包如torch、flash-attn、vllm与 CUDA/ROCm 版本存在系统级依赖——例如cu128索引上的triton并不依赖 CUDA 12.8因此它被排除在“系统依赖”之外。这与文档中“uv 会为所有相关包使用同一索引同时继续尊重生态外包的既有配置”的描述一一对应。4. 参数在 pip 接口中的接线。--torch-backend在 uv-cli 中定义为PipCompileArgs的参数并绑定环境变量UV_TORCH_BACKEND从源码结构看pip 系列的多个子命令pip compile、pip sync、pip install见 lib.rs 中的 L2209、L2561都携带该参数与文档“目前仅在uv pip接口可用”的说明一致。参数帮助文本还指出设置该选项后uv 会忽略PyTorch 生态包已配置的索引 URL转而使用所选后端——即它优先于tool.uv.index配置这正是自动选择与手工配置两套机制的优先级关系。小结与延伸阅读默认不配置索引时PyTorch 从 PyPI 安装Windows/macOS 得到 CPU 构建Linux 得到 CUDA 构建需要统一变体时用[[tool.uv.index]]explicit true[tool.uv.sources]把torch、torchvision以及 ROCm/XPU 场景下的 Triton 包钉到指定索引CUDA/ROCm/XPU 记得按平台加sys_platformmarker需要按平台分流时用环境标记需要按用户选择切换时用 optional dependencies 加conflicts声明互斥GPU 扩展包flash-attn、vllm等可用 Astral GPU 索引的预构建 wheel免去本地 CUDA 编译在uv pip接口下可用--torch-backendauto或UV_TORCH_BACKEND让 uv 根据本机 GPU 自动选索引。进一步阅读本文档原型docs/guides/integration/pytorch.md自动后端选择实现crates/uv-torch/src/backend.rs、crates/uv-torch/src/accelerator.rs参数定义与说明crates/uv-cli/src/lib.rs索引与 sources 的 schema 定义见 uv.schema.json【免费下载链接】uvAn extremely fast Python package and project manager, written in Rust.项目地址: https://gitcode.com/GitHub_Trending/uv/uv创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考