Related Resources
AI 平台/部署

利用 Wheel Variant 简化 CUDA 加速 Python 的安装和打包工作流程

2025年 8月 13日
By

如果您曾经安装过 NVIDIA GPU 加速的 Python 软件包，您可能遇到过这样的场景：导航到 pytorch.org、jax.dev、rapids.ai 或类似网站，以查找针对您的 NVIDIA CUDA 版本构建的工件。然后，您复制一个自定义 pip、uv 或其他安装程序命令，其中包含特殊的索引 URL 或特殊的软件包名称，例如 nvidia-<package>-cu{11-12}。这不仅带来不便，而且从根本上限制了 Python 软件包在现代计算环境中处理硬件多样性的方式。当前的 wheel 格式是为 CPU 计算和相对同质的计算而设计的，无法应对当今异构计算的现实。为了解决这个问题（以及其他一些问题），NVIDIA 发起了 WheelNext 开源计划。该计划旨在改善 Python 软件包生态系统的用户体验，以更好地解决科学计算、AI 和高性能计算 (HPC) 的用例。它代表了开源领域的一项重大承诺，即发展和改进许多人所依赖的 Python 生态系统。查看 WheelNext GitHub 资源库。

NVIDIA 与 Meta、Astral 和 Quansight 合作，在 PyTorch 2.8.0 中发布了对 Wheel Variant（一种新开发的格式）的实验性支持。这种新格式使您能够非常精确地描述 Python 工件，并在安装时决定哪种工件最适合该平台。本文提供了一项功能预览，解释了拟议的更改以及它们在现实世界中的工作原理。

CUDA 兼容性存在哪些技术挑战？

Python wheel 格式使用标签来识别兼容的平台：Python 版本、ABI 和平台，例如 cp313-cp313-linux_x86_64。虽然这些标签适用于基于 CPU 的软件包，但它们缺乏专门构建所需的粒度：GPU 或特定 CPU 指令 AVX512、ARMv9 等。例如，单个 linux_x86_64 标签无法说明运行 GPU 支持的软件包所需的任何其他硬件。这种粒度差距迫使软件包维护者采用次优的部署策略。

不同 CUDA 组件之间的关系通常被误解，这增加了这种复杂性，往往会导致额外的困难。这些组件包括：

  • 内核模式驱动程序 (KMD) ：将 NVIDIA 硬件与操作系统内核连接的低级固件驱动程序 (Linux 上的 nvidia.ko) 。
  • CUDA 用户模式驱动程序 (UMD) ：用于在 NVIDIA GPU 上编写和运行代码的用户模式驱动程序 (Linux 上的 libcuda.so) 。
  • CUDA Runtime：大多数 CUDA 库和应用程序使用的面向高级用户的 API（例如 cudaMemcpy）（Linux 上的 libcudart.so）。
  • CUDA Toolkit：完整的开发环境，包括编译器、库和工具。

每个组件都有不同的兼容性规则，这些规则是分布式问题的核心。

什么是 Wheel Variant 格式？

Wheel Variant 格式是一种即将提出的 Python 打包标准，旨在为异构计算时代发展 Python 打包。Wheel Variants 扩展了当前的 Wheel 格式，以便为同一软件包版本、Python ABI 和平台提供多个 Wheel，每个 Wheel 都针对特定的硬件配置进行了优化。

尝试以下操作：

# Linux
curl -LsSf https://astral.sh/uv/install.sh | INSTALLER_DOWNLOAD_URL=https://wheelnext.astral.sh sh

# Windows
powershell -c { $env:INSTALLER_DOWNLOAD_URL = 'https://wheelnext.astral.sh'; irm https://astral.sh/uv/install.ps1 | iex }

uv pip install torch

由科学计算 Python 社区开发并面向科学计算 Python 社区的协作解决方案

Python Artifact 现在可以针对特定硬件进行专门定制，从而显著改善最终用户体验和性能，而不是专注于最低公分母。

Wheel Variant 设计提出了一种优雅、简单但功能强大的语法，用于指定、识别和描述每个工件，这些工件具有遵循标准化格式的“变异属性”：

namespace :: feature :: value

变异属性的一些示例包括：

  • nvidia :: cuda_version_lower_bound :: 12.0 指定 CUDA 用户模式驱动程序 >=12.0
  • nvidia :: sm_arch :: 100_real 指定了使用 NVIDIA GPU “real architecture 100” 的 CMAKE 标记构建的软件包
  • x86_64 :: level :: v3 指定 x86-64-v3 CPU 架构支持
  • x86_64 :: avx512_bf16 :: 1 指定使用 X86_64 指令：AVX512-BF16
  • aarch64 :: version :: 8.1a 指定 ARM 架构版本 8.1a

然后，使用在构建时手动提供的自定义标签或变异属性的自动生成的 SHA-256 哈希值，对每个变异体或配置进行唯一标识。

然后，标签或哈希值将以以下方式合并到 Wheel 文件名中：

torch-2.8.0-cp313-cp313-linux_x86_64-cu128.whl     # Custom label
torch-2.8.0-cp313-cp313-linux_x86_64-a7f3c2d9.whl  # Hash-based

此设计可确保以下属性：

  • 通过提供唯一标识符，避免相同 Python 平台出现名称冲突。
  • 可选的人类可读标签，可以突出显示变异体的预期用途（例如，cu128 表示为 CUDA 12.8 或更高版本而构建）。
  • 它永远不会匹配之前的 Wheel Regex，这保证了变异轮不会混淆非变异启用的 Python 软件包安装程序。

插件架构如何工作？

通过提供商插件实现这一功能，这些插件是专门用于检测本地软件和硬件功能和配置的模块。它们分析本地环境并指导软件包选择。

运行 [uv] pip install torch 时，安装程序会查询已安装的插件，以了解系统的功能。

不同的变异插件可能会检测到您有：

  • 安装 CUDA 驱动程序 12.9
  • NVIDIA RTX 4090 GPU（计算能力 8.9）
  • 对特定 CPU 指令的支持：（例如 AVX512-BF16）

根据这些信息，安装程序会自动选择最佳 Wheel 变体：

  • 无需手动选择 CUDA 版本
  • 不再下载错误的 PyTorch 版本
  • 无需猜测；自动安装最合适的软件包
Flow chart showing (left to right) User runs > Plugin detects > Installer selects.
图 1。轮式变异体安装工作流

重要的是，Wheel Variants 保持完全向后兼容性。无法识别变异株的旧版 pip 只是忽略它们，确保现有基础设施继续运行。元数据位于三个位置：

  • pyproject.toml 用于构建配置
  • Wheel 内的 variant.json
  • *-variants.json：用于高效变异检测的软件包索引

这种设计允许逐步采用生态系统，而不会造成破坏性变化。

NVIDIA GPU 具体实现示例

NVIDIA 变量插件实施了一个优先级系统，用于处理基于 GPU 软件包安装中常见痛点的 GPU 环境的复杂性：

  • 优先级 1 (P1) – libcuda (用户模式驱动：UMD) 版本检测：这是最关键的功能。UMD 版本决定了可以使用的 CUDA Runtime 版本。如今，安装失败最常见的原因就是这些不匹配。
  • 优先级 2 (P2) – 计算能力：确定驱动程序是否包含与驱动程序安装系统上的 GPU 架构兼容的二进制代码。

举个例子，假设 NVIDIA GPU 用户使用 CUDA 驱动程序 12.8，运行：

[uv] pip install torch

幕后发生了什么？

  1. NVIDIA 插件可检测驱动程序版本和计算能力。
  2. 它查询 PyTorch 的可用变体。
  3. 它可以找到以下变异株： torch-2.8.0-...-00000000.whl – 仅 CPU 构建 – 如果其他任何内容都不匹配，则回退：称为 null-variant。 torch-2.8.0-...-cu126.whl – CTK 12.6，兼容 LibCUDA 12.0 及以上版本 torch-2.8.0-...-cu128.whl – CTK 12.8，兼容 LibCUDA 12.8 及以上版本 torch-2.8.0-...-cu129.whl – CTK 12.9，兼容 LibCUDA 12.9 及以上版本
  4. 它会选择 CUDA 12.8 变异体，因为它最匹配，然后下载并安装它。

该系统能够优雅地处理边缘情况。如果未检测到 CUDA 环境，则可以回退到 null 变体。

Python 软件包最终用户可以获得哪些生态系统优势？

Python 软件包用户可以立即获得实质性的优势：

  • 零配置安装：无需访问 pytorch.org（或 RAPIDS、JAX 等其他软件包选择器）来找出最匹配的软件包。只需使用 [uv] pip install <package> 即可。
  • 默认情况下，性能达到最佳：您会自动获得针对特定硬件编译的代码，而不是最低公分母构建。
  • 必要时用户控制：高级用户仍可使用 --no-variant 覆盖选择，以强制使用通用轮毂，或使用 <package>#<label_name> 选择特定变体。

软件包维护者可以获得哪些生态系统优势？

Wheel Variants 解决了软件包维护者长期以来遇到的分发问题，包括：

  • 简化的发布矩阵：维护者发布单个软件包的变体，而不是维护单独的软件包（torch-cpu、torch-cu126、torch-cu128、torch-cu129 等）或多个索引。
  • 有针对性的优化：针对特定架构构建经过优化的变体，而无需担心会增加其他人的下载量。
  • 面向未来的架构：当 NVIDIA 发布新的 GPU 架构时，您可以添加新的变体，而不会影响现有用户。
  • 减少支持负担：安装问题减少意味着帮助用户解决 CUDA 兼容性问题的时间减少。

NVIDIA GPU 用户可以获得哪些生态系统优势？

Wheel 变体解决了几个战略性用户体验问题，包括：

  • 解决依赖链：系统可以表达特定于变异株的依赖关系。例如，torch 可以声明特定于变异体的依赖项，从而保证所有构建中的元数据保持一致。许多软件包安装程序（例如 uv 和 poetry）都希望拥有此功能。
  • 生态系统一致性：所有 NVIDIA 库（cuDNN、cuBLAS、NCCL）均可使用一致的变异方案，使整个堆栈更具可预测性。
  • 创新支持：无需等待整个生态系统的协调，即可通过新变体公开新的 NVIDIA GPU 功能。

除了 PyTorch 和 NVIDIA GPU 计算之外，Wheel Variants 的更广泛的应用是什么？

Wheel Variants 的影响远远超出了 PyTorch 和 NVIDIA GPU 计算的范畴。下面详细介绍了几个实际应用。

研究计算环境

在学术环境中，NVIDIA GPU 集群通常具有异构硬件。一些节点使用 NVIDIA A100 GPU，其他节点使用 NVIDIA H100 GPU，较新的节点使用 NVIDIA GB200 GPU。如今，研究人员必须仔细管理环境模块和安装脚本。借助 Wheel Variants：

# On any node, regardless of NVIDIA GPU available
# Results: Automatically gets A100, H100, or B200 optimized and compatible builds
uv pip install torch transformers accelerate

系统缓存可确保平台检测只执行一次，而不是在每次安装时执行。


轮毂尺寸优化

用于 ML 和 AI 工作负载的 Docker 镜像通常很大，通常超过 10 GB。其中大部分来自适用于所有可能架构的捆绑式 GPU 库。

借助 Wheel Variants，未来可能会出现分片轮（如果软件包维护者决定分片其软件包）。

# Build stage can target specific variants
FROM nvidia/cuda:12.9.1-base-ubuntu24.04

# Sharded wheels could lead to significantly smaller docker images
RUN uv pip install torch#sm120  # Select label `sm120

对于部署数百个容器实例的组织，带宽和存储节省量迅速增加。

解决 Python 生态系统中的问题

虽然本文重点介绍了 NVIDIA GPU 的用例，但 Wheel Variants 能够解决整个 Python 生态系统中的问题，包括：

  • SciPy 可以提供针对不同 BLAS 实现（OpenBLAS 与 Intel MKL 与 Apple Accelerate）构建的变体。
  • Web assembly 项目可以提供有或无线程支持的变体（wasm32-wasi 和 pthreads）。
  • mpi4py 等 MPI 应用可以针对特定的 MPI 实现（OpenMPI 与 MPICH 与 Intel MPI）。
  • 游戏开发库可以提供具有不同图形后端（Vulkan、DirectX 或 Metal）的变体。

该变异系统可扩展到任何可以通过变异语义表达的硬件或软件功能：

namespace :: feature :: value

如何构建 Wheel 变体

本节将介绍项目维护者的用户故事，并重点介绍构建和发布 Wheel 变体的预期工作流。有许多可能的策略，从使用重新打包 CLI 工具到直接集成到大多数主要构建前端和后端。

将标准 Wheel 转换为 Wheel 变体

为了促进转换，variantlib 提供了一个 CLI 工具，用于将现有工件转换为 Wheel Variant。能够在不部署整个新流程的情况下，立即发布实验性 Wheel Variant 软件包，这一点非常有用。如需了解更多详情，请参阅 wheelnext/variantlib GitHub 资源库。

variantlib make-variant \
   -f "torch-2.8.0-cp313-cp313-manylinux_2_28_x86_64.whl" \
   -o "output_dir/" \
   --pyproject_toml "torch/pyproject.toml" \
   --variant-label "cu129" \
   --property "nvidia :: cuda_version_lower_bound :: 12.9"

Flit 和 Flit-Core：用于加速数据中心网络的软件和硬件

Flit 不支持构建 Python C 扩展；但是，它可以用于打包预构建工件。如需了解更多详情，请参阅 wheelnext/flit GitHub 资源库。

# Build a variant for x86-64-v3 architecture
flit build --format wheel \
   --variant-property "x86_64 :: level :: v3" \
   --variant-label "x8664_v3"

# Build a variant for ARM AArch64 8.1a architecture
flit build --format wheel \
   --variant-property "aarch64 :: version :: 8.1a" \
   --variant-label "arm_81a"

# Build a CUDA &gt;=12.8 variant with custom label
flit build --format wheel \
   --variant-property "nvidia :: cuda_version_lower_bound :: 12.8" \
   --variant-label "cu128"
# Build a CUDA &gt;=12.0,&lt;13 variant with custom label
flit build --format wheel \
   --variant-property "nvidia :: cuda_version_lower_bound :: 12" \
   --variant-property "nvidia :: cuda_version_upper_bound :: 13" \
   --variant-label "cu12"

# Build a null variant (no variant match =&gt; fallback)
flit build --format wheel --null-variant

Hatch 和 Hatchling

Hatch 和 Hatchling 也不支持构建 Python C 扩展，但可用于打包预构建工件。如需了解更多详情，请参阅 wheelnext/hatch GitHub 资源库。

# Build a variant for x86-64-v3 architecture
hatch build --target wheel \
   --variant-property "x86_64 :: level :: v3" \
   --variant-label "x8664_v3"

# Build a variant for ARM AArch64 8.1a architecture
hatch build --target wheel \
   --variant-property "aarch64 :: version :: 8.1a" \
   --variant-label "arm_81a"

# Build a CUDA &gt;=12.8 variant with custom label
hatch build --target wheel \
   --variant-property "nvidia :: cuda_version_lower_bound :: 12.8" \
   --variant-label "cu128"
# Build a CUDA &gt;=12.0,&lt;13 variant with custom label
hatch build --target wheel \
   --variant-property "nvidia :: cuda_version_lower_bound :: 12" \
   --variant-property "nvidia :: cuda_version_upper_bound :: 13" \
   --variant-label "cu12"

# Build a null variant (no variant match =&gt; fallback)
hatch build --target wheel --null-variant

Meson-Python：用于构建和部署 Python 程序的工具。

软件包维护者可以使用经过修改的 meson-python 构建后端开始尝试使用变体：

# Build a variant for x86-64-v3 architecture
python -m build -w -Cvariant="x86_64 :: level :: v3" -Cvariant-label=x8664_v3

# Build a variant for ARM AArch64 8.1a architecture
python -m build -w -Cvariant="aarch64 :: version :: 8.1a" -Cvariant-label=arm_81a

# Build a CUDA 12.8+ variant with custom label
python -m build -w -Cvariant="nvidia :: cuda_version_lower_bound :: 12.8" -Cvariant-label=cu128

# Build a null variant (no variant match =&gt; fallback)
python -m build -w -Cnull-variant

构建系统将变体信息传递给编译过程，从而在保持单个源代码树的同时实现有针对性的优化。如需了解更多详情，请参阅 wheelnext/meson-python GitHub 资源库。

Scikit-Build-Core：用于构建和训练模型的工具。

这项工作正在进行中，敬请关注更新和更多信息。

Wheel 变异体的实施路线图是什么？

通过细致的协作开发，轮毂变异计划正在从概念走向现实。

  • PyTorch 2.8.0：将包括对 Wheel Variants 的实验性支持。这是 PyTorch 团队与开源 WheelNext 项目、Quansight、Astral 和 NVIDIA 合作发布的测试版本。这项实验将使团队能够在 Python 增强建议 (PEP) 发布之前收集现实世界的反馈。
  • PEP 处于草案状态：社区审核和改进确保所有利益相关者都能为塑造标准做出贡献。
  • 参考实现：variantlib 库提供核心功能，而原型 uv 和 pip 实现展示了操作模型以及与现有工具和生态系统的共存。
  • NVIDIA 插件正在积极开发中：实现本文中描述的优先级系统。

展望未来，路线图优先考虑生态系统兼容性和逐步采用。

  • 近期：在 PyTorch 2.8 中进行实验阶段，收集反馈意见
  • 中期：基于实际经验完成 PEP 终稿
  • 长期：更广泛的工具支持（安装程序、构建后端和前端、软件包索引）和整个生态系统的采用

重要的是，这是一项协作性工作。成功取决于科学 Python 社区的各个角落的共同努力，他们认识到该标准的价值并采用该标准。

结论

Wheel 变体不仅仅代表了技术上的改进。它们是 Python 软件包处理日益多样化的计算环境的基本演变。对于 NVIDIA 生态系统而言，它们解决了当前的痛点，同时将 GPU 计算确立为 Python 打包的基本方面。

PyTorch、Astral、Quansight、NVIDIA 和 WheelNext 项目之间的合作展示了 Python 社区的最佳状态：识别共同的挑战并构建使每个人受益的解决方案。PyTorch 2.8 测试版是 Python 打包新时代的开端，在这个时代，[uv] pip install <package> 能够以优化方式适用于您的本地硬件和软件配置。

我们邀请您加入我们的旅程。测试实验版本、提供反馈并帮助塑造 Python 打包的未来。我们携手合作，确保下一代 Python 开发者再也不用手动选择 CUDA 版本了。

如需了解更多信息并参与其中，请访问 wheelnext.dev，并关注 PyTorch 2.8 实验性 Wheel 变异体支持。

我们鼓励您阅读我们的合作者的博客文章，进一步了解 Wheel 变体提案：

参与其中

Wheel Variants 计划需要社区参与才能取得成功。以下是不同群体可以做出的贡献：

开发者和用户：

  • 在 PyTorch 2.8 发布时测试其实验性支持。现实世界的反馈非常宝贵。
  • 报告您遇到的问题和边缘情况。系统必须能够处理各种环境。
  • 分享您的用例，这些用例可以从变异体中受益。

软件包维护者：

  • 探索各种包装方案。wheelnext.dev 网站、Wheel 变体提案和 WheelNext GitHub 网站提供了全面的文档。
  • 考虑您的硬件矩阵。哪些变体对您的用户有益？
  • 提供对 API 的反馈。PEP 仍处于草案状态。您的输入将塑造最终提案。

欢迎参加 PyTorch Conference 2025，观看题为“硬件感知 Python 软件包 ~ PyTorch 和 WheelNext Grab the Wheel”的演讲。

致谢

这项工作之所以能够完成，要归功于许多开源贡献者和 WheelNext 社区成员的努力。其中包括：Michał Górny (Quansight) 、Ralf Gommers (Quansight) 、Charlie Marsh (Astral) 、Konstantin Schütze (Astral) 、Zanie Blue (Astral) 、Andrey Talman (Meta) 、Eli Uriegas (Meta) 、Chris Gottbrath (Meta) 以及许多 NVIDIA 员工：Andy Terrel、Barry Warsaw、Emma Smith、Michael Sarahan、Vyas Ramasubramani、Bradley Dice、Robert Maynard、Hyunsu Cho、Ralf W. Grosse Kunstleve、Leo Fang、Keith Kraus、Piotr Bialecki、Frederic Bastien、Jeremy Tanner、David Edelsohn、Jay Gould、Scott Suchyta、Ankit Patel。

 

相关资源

标签

AI 平台/部署 | 开发与优化 | 模型/库/框架 | AI Platforms / Deployment | CUDA | Development & Optimization | HPC/科学计算 | Intermediate Technical | Models / Libraries / Frameworks | open source | Python | PyTorch | RTX GPU | 最佳实践

关于作者

自 2019 年以来，Jonathan Dekhtiar 一直在为许多深度学习库、编译器和框架做出贡献。在 NVIDIA，他在设计 Python 软件和库发布机制以及解决 GPU 加速 Python 代码的打包挑战方面发挥了关键作用。作为 Wheel-Next 计划的推动者和创始者，Jonathan 投入了大量精力与社区一起推进这一开源工作，并积极致力于概念验证、演示和 PEP 的编写。